beam: O(1) LFU cache with doubly-linked freq buckets + tests

wizard-work/beam/lfu_cache.py

@@ -0,0 +1,108 @@
+"""O(1) Least-Frequently-Used cache.
+
+Design: each frequency bucket is a doubly-linked list ordered by recency
+(most-recently-used at the head). A node stores key, value, freq. On access
+the node moves from its current bucket to the next-frequency bucket's head.
+min_freq tracks the bucket containing the least-frequent (and among those,
+least-recently-used) node for O(1) eviction.
+"""
+
+
+class _Node:
+    __slots__ = ("key", "value", "freq", "prev", "next")
+
+    def __init__(self, key, value):
+        self.key = key
+        self.value = value
+        self.freq = 1
+        self.prev = None
+        self.next = None
+
+
+class _DLList:
+    """Doubly linked list with sentinel head/tail; most-recent at head side."""
+
+    __slots__ = ("head", "tail", "size")
+
+    def __init__(self):
+        self.head = _Node(None, None)  # sentinel
+        self.tail = _Node(None, None)  # sentinel
+        self.head.next = self.tail
+        self.tail.prev = self.head
+        self.size = 0
+
+    def add_front(self, node):
+        node.prev = self.head
+        node.next = self.head.next
+        self.head.next.prev = node
+        self.head.next = node
+        self.size += 1
+
+    def remove(self, node):
+        node.prev.next = node.next
+        node.next.prev = node.prev
+        node.prev = None
+        node.next = None
+        self.size -= 1
+
+    def pop_tail(self):
+        # least recent = node before tail sentinel
+        if self.size == 0:
+            return None
+        node = self.tail.prev
+        self.remove(node)
+        return node
+
+    def is_empty(self):
+        return self.size == 0
+
+
+class LFUCache:
+    def __init__(self, capacity):
+        if capacity < 0:
+            raise ValueError("capacity must be non-negative")
+        self.capacity = capacity
+        self._key = {}          # key -> _Node
+        self._freq = {}         # freq -> _DLList
+        self._min_freq = 0
+
+    def _touch(self, node):
+        # move node from its current freq bucket to freq+1 bucket head
+        old_list = self._freq[node.freq]
+        old_list.remove(node)
+        if old_list.is_empty():
+            del self._freq[node.freq]
+            if self._min_freq == node.freq:
+                self._min_freq += 1
+        node.freq += 1
+        self._freq.setdefault(node.freq, _DLList()).add_front(node)
+
+    def get(self, key):
+        if key not in self._key:
+            return -1
+        node = self._key[key]
+        self._touch(node)
+        return node.value
+
+    def put(self, key, value):
+        if self.capacity == 0:
+            return
+        if key in self._key:
+            node = self._key[key]
+            node.value = value
+            self._touch(node)
+            return
+        # eviction if needed
+        if len(self._key) >= self.capacity:
+            evict_list = self._freq.get(self._min_freq)
+            if evict_list is not None:
+                node = evict_list.pop_tail()
+                if node is not None:
+                    del self._key[node.key]
+                    if evict_list.is_empty():
+                        del self._freq[self._min_freq]
+        # insert new node at freq 1
+        node = _Node(key, value)
+        self._key[key] = node
+        self._freq.setdefault(1, _DLList()).add_front(node)
+        self._min_freq = 1

wizard-work/beam/test_lfu_cache.py

@@ -0,0 +1,133 @@
+import sys
+sys.path.insert(0, "/root/.hermes/wizard-outputs/beam")
+from lfu_cache import LFUCache
+
+
+def test_basic_get_put():
+    c = LFUCache(2)
+    c.put(1, 10)
+    c.put(2, 20)
+    assert c.get(1) == 10
+    assert c.get(2) == 20
+    assert c.get(3) == -1
+
+
+def test_evict_least_freq():
+    c = LFUCache(2)
+    c.put(1, 10)
+    c.put(2, 20)
+    c.get(1)  # freq 1->2
+    c.put(3, 30)  # evict 2 (freq1)
+    assert c.get(2) == -1
+    assert c.get(1) == 10
+    assert c.get(3) == 30
+
+
+def test_lru_within_same_freq():
+    c = LFUCache(2)
+    c.put(1, 10)
+    c.put(2, 20)
+    # both freq1; touch nothing extra; insert 3 evicts LRU among freq1 -> key1
+    c.put(3, 30)
+    assert c.get(1) == -1
+    assert c.get(2) == 20
+    assert c.get(3) == 30
+
+
+def test_update_existing():
+    c = LFUCache(2)
+    c.put(1, 10)
+    c.put(1, 11)
+    assert c.get(1) == 11
+    # updating counts as access; freq now 2
+
+
+def test_capacity_zero():
+    c = LFUCache(0)
+    c.put(1, 10)
+    assert c.get(1) == -1
+
+
+def test_evict_after_freq_tiebreak():
+    c = LFUCache(3)
+    c.put(1, 1)
+    c.put(2, 2)
+    c.put(3, 3)
+    c.get(1)
+    c.get(2)
+    # freq: 1->2, 2->2, 3->1 ; min freq bucket {3}
+    c.put(4, 4)  # evict 3
+    assert c.get(3) == -1
+    assert c.get(1) == 1 and c.get(2) == 2 and c.get(4) == 4
+
+
+def test_negative_capacity_raises():
+    try:
+        LFUCache(-1)
+        assert False
+    except ValueError:
+        pass
+
+
+def test_many_ops_stress():
+    # brute-force compare against a naive dict-based LFU
+    import random
+    random.seed(42)
+    cap = 5
+    c = LFUCache(cap)
+    # naive
+    store = {}
+    freq = {}
+    order = {}  # key -> insertion/access counter for LRU tiebreak
+    counter = [0]
+
+    def naive_get(k):
+        if k not in store:
+            return -1
+        freq[k] += 1
+        counter[0] += 1
+        order[k] = counter[0]
+        return store[k]
+
+    def naive_put(k, v):
+        nonlocal store
+        if cap == 0:
+            return
+        if k in store:
+            store[k] = v
+            freq[k] += 1
+            counter[0] += 1
+            order[k] = counter[0]
+            return
+        if len(store) >= cap:
+            # evict min freq, then min order
+            minf = min(freq[x] for x in store)
+            candidates = [x for x in store if freq[x] == minf]
+            victim = min(candidates, key=lambda x: order[x])
+            del store[victim]; del freq[victim]; del order[victim]
+        store[k] = v
+        freq[k] = 1
+        counter[0] += 1
+        order[k] = counter[0]
+
+    ops = 2000
+    keys = list(range(1, 12))
+    for _ in range(ops):
+        k = random.choice(keys)
+        if random.random() < 0.5:
+            a = c.get(k)
+            b = naive_get(k)
+            assert a == b, (a, b, k)
+        else:
+            v = random.randint(0, 99)
+            c.put(k, v)
+            naive_put(k, v)
+    # final consistency: all keys
+    for k in keys:
+        assert c.get(k) == naive_get(k), k
+
+
+if __name__ == "__main__":
+    for name in [n for n in globals() if n.startswith("test_")]:
+        globals()[name]()
+    print("ALL TESTS PASSED")