105 lines
3.3 KiB
Python
105 lines
3.3 KiB
Python
"""Thread-safe O(1) LRU cache (doubly linked list + hash map).
|
|
|
|
Subtle bugs fixed vs. naive versions:
|
|
- `get`/`put`-on-existing BOTH promote to MRU (forgetting promote on get is the
|
|
classic LRU bug).
|
|
- Eviction deletes from `_map` (forgetting this leaks memory and breaks size).
|
|
- Missing-key sentinel uses a unique `_MISSING` object instead of `-1` so a
|
|
stored value of -1 is not mistaken for "absent".
|
|
- All public ops hold a single RLock so concurrent access cannot corrupt the
|
|
linked list or hand out a value for an evicted key.
|
|
- Capacity validated as positive int (bool rejected).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
from typing import Any, Optional
|
|
|
|
_MISSING = object()
|
|
|
|
|
|
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 isinstance(capacity, bool) or not isinstance(capacity, int):
|
|
raise TypeError("capacity must be an int")
|
|
if capacity <= 0:
|
|
raise ValueError("capacity must be positive")
|
|
self.capacity = capacity
|
|
self._map: dict[Any, _Node] = {}
|
|
self._head = _Node(None, None) # MRU side sentinel
|
|
self._tail = _Node(None, None) # LRU side sentinel
|
|
self._head.next = self._tail
|
|
self._tail.prev = self._head
|
|
self._lock = threading.RLock()
|
|
|
|
def _detach(self, node: _Node) -> None:
|
|
p, n = node.prev, 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:
|
|
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)
|
|
|
|
def get(self, key: Any, default: Any = _MISSING) -> Any:
|
|
with self._lock:
|
|
node = self._map.get(key)
|
|
if node is None:
|
|
return None if default is _MISSING else default
|
|
self._promote(node)
|
|
return node.value
|
|
|
|
def put(self, key: Any, value: Any) -> None:
|
|
with self._lock:
|
|
node = self._map.get(key)
|
|
if node is not None:
|
|
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:
|
|
lru = self._tail.prev
|
|
assert lru is not None and lru is not self._head
|
|
self._detach(lru)
|
|
del self._map[lru.key]
|
|
|
|
def contains(self, key: Any) -> bool:
|
|
with self._lock:
|
|
return key in self._map
|
|
|
|
def size(self) -> int:
|
|
with self._lock:
|
|
return len(self._map)
|
|
|
|
def order_mru_to_lru(self) -> list[Any]:
|
|
with self._lock:
|
|
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
|