beam: O(1) LRU cache with randomized property tests

wizard-work/beam/lru_cache.py

@@ -0,0 +1,98 @@
+"""O(1) LRU cache: doubly linked list + hash map.
+
+Design notes:
+- Each node holds key, value, prev, next.
+- `_map` maps key -> node.
+- Head is MRU sentinel, tail is LRU sentinel. Most-recently-used sits right after head.
+- `get` and `put`-on-existing must BOTH promote the node to MRU.
+- `put`-on-new inserts at MRU; if over capacity, evicts the node before tail (LRU).
+- Eviction must delete from `_map` (common bug: forget this).
+"""
+from __future__ import annotations
+
+from typing import Any, Optional
+
+
+class _Node:
+    __slots__ = ("key", "value", "prev", "next")
+
+    def __init__(self, key: Any, value: Any) -> None:
+        self.key = key
+        self.value = value
+        self.prev: Optional["_Node"] = None
+        self.next: Optional["_Node"] = None
+
+
+class LRUCache:
+    def __init__(self, capacity: int) -> None:
+        if not isinstance(capacity, int) or isinstance(capacity, bool):
+            raise TypeError("capacity must be an int")
+        if capacity <= 0:
+            raise ValueError("capacity must be positive")
+        self.capacity = capacity
+        self._map: dict[Any, _Node] = {}
+        # sentinels
+        self._head = _Node(None, None)  # MRU side
+        self._tail = _Node(None, None)  # LRU side
+        self._head.next = self._tail
+        self._tail.prev = self._head
+
+    # ---- internal list ops ----
+    def _detach(self, node: _Node) -> None:
+        p = node.prev
+        n = node.next
+        assert p is not None and n is not None
+        p.next = n
+        n.prev = p
+        node.prev = None
+        node.next = None
+
+    def _attach_mru(self, node: _Node) -> None:
+        # insert right after head
+        first = self._head.next
+        self._head.next = node
+        node.prev = self._head
+        node.next = first
+        assert first is not None
+        first.prev = node
+
+    def _promote(self, node: _Node) -> None:
+        self._detach(node)
+        self._attach_mru(node)
+
+    # ---- public API ----
+    def get(self, key: Any) -> Any:
+        node = self._map.get(key)
+        if node is None:
+            return -1
+        self._promote(node)
+        return node.value
+
+    def put(self, key: Any, value: Any) -> None:
+        node = self._map.get(key)
+        if node is not None:
+            # update value AND recency (subtle bug if you skip promote)
+            node.value = value
+            self._promote(node)
+            return
+        node = _Node(key, value)
+        self._map[key] = node
+        self._attach_mru(node)
+        if len(self._map) > self.capacity:
+            # evict LRU: node just before tail
+            lru = self._tail.prev
+            assert lru is not None and lru is not self._head
+            self._detach(lru)
+            del self._map[lru.key]
+
+    def size(self) -> int:
+        return len(self._map)
+
+    # helpers for tests / introspection
+    def order_mru_to_lru(self) -> list[Any]:
+        out = []
+        cur = self._head.next
+        while cur is not None and cur is not self._tail:
+            out.append(cur.key)
+            cur = cur.next
+        return out

wizard-work/beam/test_lru_cache.py

@@ -0,0 +1,84 @@
+import random
+from collections import OrderedDict
+
+from lru_cache import LRUCache
+
+
+def test_basic_get_put():
+    c = LRUCache(2)
+    c.put(1, 1)
+    c.put(2, 2)
+    assert c.get(1) == 1
+    c.put(3, 3)  # evicts 2
+    assert c.get(2) == -1
+    assert c.get(1) == 1
+    assert c.get(3) == 3
+
+
+def test_update_promotes_recency():
+    # subtle bug: updating existing key must move to MRU
+    c = LRUCache(2)
+    c.put(1, 1)
+    c.put(2, 2)
+    c.put(1, 10)  # update -> 1 becomes MRU, 2 LRU
+    c.put(3, 3)   # should evict 2, not 1
+    assert c.get(2) == -1
+    assert c.get(1) == 10
+    assert c.get(3) == 3
+
+
+def test_eviction_removes_from_map():
+    c = LRUCache(1)
+    c.put(1, 1)
+    c.put(2, 2)
+    assert c.size() == 1
+    assert c.get(1) == -1
+
+
+def test_capacity_validation():
+    for bad in (0, -1, True, "2", 2.0):
+        try:
+            LRUCache(bad)
+            raised = False
+        except (ValueError, TypeError):
+            raised = True
+        assert raised, bad
+
+
+def test_order_introspection():
+    c = LRUCache(3)
+    for k in (1, 2, 3):
+        c.put(k, k)
+    assert c.order_mru_to_lru() == [3, 2, 1]
+    c.get(1)
+    assert c.order_mru_to_lru() == [1, 3, 2]
+
+
+def test_randomized_against_ordereddict():
+    random.seed(42)
+    cap = 5
+    c = LRUCache(cap)
+    ref = OrderedDict()
+    ops = 2000
+    for _ in range(ops):
+        k = random.randrange(0, 12)
+        if random.random() < 0.5:
+            v = random.randrange(0, 100)
+            c.put(k, v)
+            ref[k] = v
+            # move to end (MRU) in reference
+            ref.move_to_end(k)
+            if len(ref) > cap:
+                ref.popitem(last=False)
+        else:
+            got = c.get(k)
+            if k in ref:
+                ref.move_to_end(k)
+                assert got == ref[k], (k, got, ref[k])
+            else:
+                assert got == -1
+        # size invariant
+        assert c.size() == len(ref)
+        # order invariant
+        expected_order = list(reversed(list(ref.keys())))  # MRU..LRU
+        assert c.order_mru_to_lru() == expected_order, (c.order_mru_to_lru(), expected_order)