modelTesting/outputs/gemma4-26b-a4b-8bit-mlx.py

import asyncio
import time
from dataclasses import dataclass, field
from typing import Any, Dict, Generic, Optional, Set, TypeVar, Union

# Type variables for Generics
K = TypeVar("K")
V = TypeVar("V")

@dataclass
class Node(Generic[K, V]):
    """A node in the doubly linked list representing a cache entry."""
    key: K
    value: V
    freq: int = 1
    expiry: float = float('inf')
    prev: Optional['Node[K, V]'] = None
    next: Optional['Node[K, V]'] = None

class DoublyLinkedList(Generic[K, V]):
    """A standard Doubly Linked List to allow O(1) removal and insertion."""
    def __init__(self):
        self.head: Optional[Node[K, V]] = None
        self.tail: Optional[Node[K, V]] = None
        self.size: int = 0

    def append(self, node: Node[K, V]):
        """Adds a node to the end of the list."""
        if not self.head:
            self.head = node
            self.tail = node
            node.prev = None
            node.next = None
        else:
            node.prev = self.tail
            node.next = None
            if self.tail:
                self.tail.next = node
            self.tail = node
        self.size += 1

    def remove(self, node: Node[K, V]):
        """Removes a specific node from the list in O(1)."""
        if node.prev:
            node.prev.next = node.next
        else:
            self.head = node.next

        if node.next:
            node.next.prev = node.prev
        else:
            self.tail = node.prev
        
        node.next = None
        node.prev = None
        self.size -= 1

    def pop_tail(self) -> Optional[Node[K, V]]:
        """Removes and returns the last node in O(1)."""
        if not self.tail:
            return None
        node = self.tail
        self.remove(node)
        return node

class Transaction(Generic[K, V]):
    """
    Implements ACID-like sub-sessions. 
    Provides 'Read Your Own Writes' and isolation from the global cache.
    """
    def __init__(self, cache: 'LFUCache[K, V]'):
        self._cache = cache
        self._pending_puts: Dict[K, tuple[V, float]] = {}
        self._pending_deletes: Set[K] = set()
        # Track keys read from the main cache to update frequency on commit
        self._read_cache_keys: Set[K] = set()
        self._is_active = True

    async def get(self, key: K) -> Optional[V]:
        if not self._is_active: raise RuntimeError("Transaction closed")
        
        # 1. Check local deletes (Tombstone)
        if key in self._pending_deletes:
            return None
        # 2. Check local writes (Read Your Own Writes)
        if key in self._pending_puts:
            val, _ = self._pending_puts[key]
            return val
        # 3. Check global cache (without affecting LFU frequency until commit)
        async with self._cache._lock:
            node = self._cache.cache_data.get(key)
            if node:
                # Lazy TTL check even during transaction get
                if time.monotonic() >= node.expiry:
                    self._cache._remove_node_from_structures(node)
                    return None
                self._read_cache_keys.add(key)
                return node.value
        return None

    async def put(self, key: K, value: V, ttl_seconds: float = float('inf')):
        if not self._is_active: raise RuntimeError("Transaction closed")
        expiry = time.monotonic() + ttl_seconds
        self._pending_puts[key] = (value, expiry)
        if key in self._pending_deletes:
            self._pending_deletes.remove(key)

    async def delete(self, key: K):
        if not self._is_active: raise RuntimeError("Transaction closed")
        if key in self._pending_puts:
            del self._pending_puts[key]
        self._pending_deletes.add(key)

    async def commit(self):
        if not self._is_active: raise RuntimeError("Transaction closed")
        async with self._cache._lock:
            # 1. Process Deletes
            for key in self._pending_deletes:
                node = self._cache.cache_data.get(key)
                if node:
                    self._cache._remove_node_from_structures(node)

            # 2. Process Reads (Update frequencies for keys read from global cache)
            for key in self._read_cache_keys:
                node = self._cache.cache_data.get(key)
                if node: # Ensure it wasn't deleted by a pending delete in this TX
                    self._cache._increment_frequency(node)

            # 3. Process Puts
            for key, (val, expiry) in self._pending_puts.items():
                # Internal un-locked put for use within the lock context of commit
                self._cache._internal_put(key, val, expiry)
        
        self._is_active = False
        self._pending_puts.clear()

    async def rollback(self):
        if not self._is_active: raise RuntimeError("Transaction closed")
        self._pending_puts.clear()
        self._pending_deletes.clear()
        self._read_cache_keys.clear()
        self._is_active = False

class LFUCache(Generic[K, V]):
    """
    In-Memory Concurrent LFU Cache.
    Time Complexity: O(1) for get and put.
    """
    def __init__(self, capacity: int):
        if capacity <= 0: raise ValueError("Capacity must be > 0")
        self.capacity = capacity
        self.cache_data: Dict[K, Node[K, V]] = {}
        self.freq_map: Dict[int, DoublyLinkedList[K, V]] = {}
        self.min_freq: int = 0
        self._lock = asyncio.Lock()
        self._evictor_task: Optional[asyncio.Task] = None

    async def start_evictor(self, interval: float = 1.0):
        """Starts the background async eviction loop."""
        if self._evictor_task is None:
            self._evictor_task = asyncio.create_task(self._background_eviction_loop(interval))

    async def stop_evictor(self):
        """Stops the background async eviction loop."""
        if self._evictor_task:
            self._evictor_task.cancel()
            try:
                await self._evictor_task
            except asyncio.CancelledError:
                pass
            self._evictor_task = None

    async def _background_eviction_loop(self, interval: float):
        while True:
            await asyncio.sleep(interval)
            # We do not hold the lock for the entire loop to prevent blocking.
            # Instead, we take a snapshot of keys and process in batches.
            async with self._lock:
                keys_snapshot = list(self.cache_data.keys())

            # Process in small batches to yield control
            batch_size = 50
            for i in range(0, len(keys_snapshot), batch_size):
                batch = keys_snapshot[i : i + batch_size]
                now = time.monotonic()
                async with self._lock:
                    for key in batch:
                        node = self.cache_data.get(key)
                        if node and now >= node.expiry:
                            self._remove_node_from_structures(node)

    def _remove_node_from_structures(self, node: Node[K, V]):
        """Internal: Removes node from freq_map and cache_data. O(1)."""
        # Remove from DLL
        dll = self.freq_map[node.freq]
        dll.remove(node)
        if dll.size == 0:
            del self.freq_map[node.freq]
            if self.min_freq == node.freq:
                # This is a simplification; real min_freq update happens in _increment_frequency.
                # If current min freq list is empty, we'll find the new min in _increment_frequency or put.
                pass 
        
        # Remove from dict
        if node.key in self.cache_data:
            del self.cache_data[node.key]

    def _increment_frequency(self, node: Node[K, V]):
        """Internal: Increases frequency of a node. O(1)."""
        old_freq = node.freq
        new_freq = old_freq + 1
        node.freq = new_freq

        # Remove from old DLL
        old_dll = self.freq_map[old_freq]
        old_dll.remove(node)
        if old_dll.size == 0:
            del self.freq_map[old_freq]
            if self.min_freq == old_freq:
                self.min_freq = new_freq

        # Add to new DLL
        if new_freq not in self.freq_map:
            self.freq_map[new_freq] = DoublyLinkedList()
        self.freq_map[new_freq].append(node)

    def _internal_put(self, key: K, value: V, expiry: float):
        """Internal: The core LFU logic. Must be called within a lock."""
        if key in self.cache_data:
            node = self.cache_data[key]
            node.value = value
            node.expiry = expiry
            self._increment_frequency(node)
        else:
            if len(self.cache_data) >= self.capacity:
                # Evict LFU (min_freq list tail)
                if self.min_freq in self.freq_map and self.freq_map[self.min_freq].size > 0:
                    victim = self.freq_map[self.min_freq].pop_tail()
                    if victim:
                        del self.cache_data[victim.key]
                else:
                    # Fallback (should not happen with correct logic)
                    k_evict = next(iter(self.cache_data))
                    del self.cache_data[k_evict]

            new_node = Node(key, value, freq=1, expiry=expiry)
            self.cache_data[key] = new_node
            if 1 not in self.freq_map:
                self.freq_map[1] = DoublyLinkedList()
            self.freq_map[1].append(new_node)
            self.min_freq = 1

    async def get(self, key: K) -> Optional[V]:
        """O(1) retrieval with lazy TTL eviction."""
        async with self._lock:
            node = self.cache_data.get(key)
            if not node:
                return None
            
            # Lazy TTL Eviction
            if time.monotonic() >= node.expiry:
                self._remove_node_from_structures(node)
                return None
            
            self._increment_frequency(node)
            return node.value

    async def put(self, key: K, value: V, ttl_seconds: float = float('inf')):
        """O(1) insertion with lazy TTL eviction."""
        expiry = time.monotonic() + ttl_seconds
        async with self._lock:
            self._internal_put(key, value, expiry)

    async def delete(self, key: K):
        """O(1) deletion."""
        async with self._lock:
            node = self.cache_data.get(key)
            if node:
                self._remove_node_from_structures(node)

    def begin_transaction(self) -> Transaction[K, V]:
        return Transaction(self)

# ==========================================
# UNIT TESTS
# ==========================================

async def test_lfu_eviction():
    print("Testing O(1) LFU Eviction Logic...")
    cache = LFUCache[str, int](capacity=3)
    await cache.put("a", 1) # freq 1
    await cache.put("b", 2) # freq 1
    await cache.put("c", 3) # freq 1
    
    # Increase frequency of a and b
    await cache.get("a") # freq 2
    await cache.get("b") # freq 2
    # c is still freq 1
    
    await cache.put("d", 4) # Should evict 'c'
    
    assert await cache.get("a") == 1
    assert await cache.get("b") == 2
    assert await cache.get("c") is None
    assert await cache.get("d") == 4
    print("✅ LFU Eviction Passed.")

async def test_ttl_eviction():
    print("Testing Dual-Layer TTL Eviction...")
    cache = LFUCache[str, int](capacity=10)
    await cache.start_evictor(interval=0.1)

    # Test Lazy Eviction
    await cache.put("lazy", 100, ttl_seconds=0.2)
    await asyncio.sleep(0.3)
    assert await cache.get("lazy") is None, "Lazy eviction failed"

    # Test Background Eviction
    await cache.put("bg", 200, ttl_seconds=0.1)
    assert await cache.get("bg") is not None, "Value should still be there for a millisecond"
    await asyncio.sleep(0.4)
    # Note: Background loop might not have run yet, but we check if it's gone
    # Since background is a separate task, it should have cleared 'bg' by now.
    assert await cache.get("bg") is None, "Background eviction failed"

    await cache.stop_evictor()
    print("✅ TTL Eviction Passed.")

async def test_transactions():
    print("Testing Atomic Transactions...")
    cache = LFUCache[str, int](capacity=5)

    # Test Rollback
    tx = cache.begin_transaction()
    await tx.put("tx1", 10)
    assert await cache.get("tx1") is None, "Uncommitted write visible!"
    assert await tx.get("tx1") == 10, "Read Your Own Writes failed"
    await tx.rollback()
    assert await cache.get("tx1") is None

    # Test Commit Visibility
    tx = cache.begin_transaction()
    await tx.put("tx2", 20)
    await tx.commit()
    assert await cache.get("tx2") == 20, "Commit visibility failed"

    # Test Isolation/Tombstones
    await cache.put("exists", 50)
    tx = cache.begin_transaction()
    await tx.delete("exists")
    assert await tx.get("exists") is None, "Transaction delete failed"
    assert await cache.get("exists") == 50, "Uncommitted delete visible!"
    await tx.commit()
    assert await cache.get("exists") is None, "Commit delete failed"

    print("✅ Transactions Passed.")

async def test_stress():
    print("Running Concurrency Stress Test (50 tasks)...")
    cache = LFUCache[int, int](capacity=20)
    await cache.start_evictor()

    async def worker(worker_id: int):
        for i in range(100):
            key = (worker_id * 100) + (i % 30) # Overlapping keys to induce contention
            op = i % 3
            if op == 0:
                await cache.put(key, i)
            elif op == 1:
                await cache.get(key)
            else:
                # Transactional stress
                tx = cache.begin_transaction()
                await tx.put(key, i)
                if i % 5 == 0:
                    await tx.rollback()
                else:
                    await tx.commit()
            if i % 10 == 0:
                await asyncio.sleep(0.01)

    tasks = [worker(i) for i in range(50)]
    await asyncio.gather(*tasks)
    await cache.stop_evictor()
    print("✅ Stress Test Completed (No crashes).")

async def main():
    start_time = time.perf_counter()
    try:
        await test_lfu_eviction()
        print("-" * 30)
        await test_ttl_eviction()
        print("-" * 30)
        await test_transactions()
        print("-" * 30)
        await test_stress()
        print("-" * 30)
        print(f"All tests finished successfully in {time.perf_counter() - start_time:.2f}s")
    except Exception as e:
        print(f"❌ Tests failed with error: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    asyncio.run(main())