runtime/cache.py

from collections import OrderedDict


class TokenCache:
    """In-memory cache for decoded token chunks.

    Frequently accessed chunks are kept hot in RAM, avoiding redundant
    entropy decode and disk reads.  LRU eviction when capacity is reached.
    """
    def __init__(self, capacity: int = 32):
        self._capacity = capacity
        self._store: OrderedDict[int, list[int]] = OrderedDict()

    def get(self, chunk_index: int) -> list[int] | None:
        if chunk_index not in self._store:
            return None
        self._store.move_to_end(chunk_index)
        return self._store[chunk_index]

    def put(self, chunk_index: int, tokens: list[int]):
        self._store[chunk_index] = tokens
        self._store.move_to_end(chunk_index)
        while len(self._store) > self._capacity:
            self._store.popitem(last=False)

    def invalidate(self, chunk_index: int):
        self._store.pop(chunk_index, None)

    def clear(self):
        self._store.clear()

    @property
    def size(self) -> int:
        return len(self._store)

    @property
    def capacity(self) -> int:
        return self._capacity