tinyquant_core/codec/
rotation_cache.rs

1//! Process-wide LRU cache for rotation matrices, keyed by `(seed, dim)`.
2//!
3//! Mirrors Python's `functools.lru_cache(maxsize=8)` on
4//! `RotationMatrix._cached_build`. The cache stores up to `capacity`
5//! entries and evicts the least-recently-used on insertion. Hits also
6//! promote the touched entry to the front of the queue.
7//!
8//! The implementation is intentionally minimal — it uses a
9//! `spin::Mutex<Vec<RotationMatrix>>` rather than a proper hash map
10//! because the expected capacity is tiny (default 8) and the linear
11//! scan fits in a single cache line at that size.
12
13use alloc::vec::Vec;
14
15use spin::Mutex;
16
17use crate::codec::rotation_matrix::RotationMatrix;
18
19/// Default cache capacity, matching Python's `functools.lru_cache(maxsize=8)`.
20pub const DEFAULT_CAPACITY: usize = 8;
21
22/// A small `(seed, dim)`-keyed LRU cache over [`RotationMatrix`] values.
23pub struct RotationCache {
24    entries: Mutex<Vec<RotationMatrix>>,
25    capacity: usize,
26}
27
28impl RotationCache {
29    /// Construct a new cache with the given capacity.
30    ///
31    /// A capacity of `0` is legal but makes every lookup a miss.
32    #[must_use]
33    pub const fn new(capacity: usize) -> Self {
34        Self {
35            entries: Mutex::new(Vec::new()),
36            capacity,
37        }
38    }
39
40    /// Construct the default-sized cache (capacity [`DEFAULT_CAPACITY`]).
41    #[must_use]
42    pub const fn default_sized() -> Self {
43        Self::new(DEFAULT_CAPACITY)
44    }
45
46    /// The maximum number of entries this cache will retain.
47    #[inline]
48    #[must_use]
49    pub const fn capacity(&self) -> usize {
50        self.capacity
51    }
52
53    /// The number of entries currently held by the cache.
54    #[must_use]
55    pub fn len(&self) -> usize {
56        self.entries.lock().len()
57    }
58
59    /// `true` iff the cache currently holds no entries.
60    #[must_use]
61    pub fn is_empty(&self) -> bool {
62        self.entries.lock().is_empty()
63    }
64
65    /// Return the rotation matrix for `(seed, dimension)`, building and
66    /// caching one on miss. Returned values are `Arc`-backed clones, so
67    /// the caller never blocks subsequent users of the cache.
68    pub fn get_or_build(&self, seed: u64, dimension: u32) -> RotationMatrix {
69        if self.capacity == 0 {
70            return RotationMatrix::build(seed, dimension);
71        }
72        let mut guard = self.entries.lock();
73        if let Some(pos) = guard
74            .iter()
75            .position(|m| m.seed() == seed && m.dimension() == dimension)
76        {
77            let touched = guard.remove(pos);
78            let clone = touched.clone();
79            guard.insert(0, touched);
80            return clone;
81        }
82        // Miss: build, evict if necessary, insert at MRU position.
83        let built = RotationMatrix::build(seed, dimension);
84        if guard.len() >= self.capacity {
85            guard.pop();
86        }
87        guard.insert(0, built.clone());
88        built
89    }
90
91    /// Drop all cached entries.
92    pub fn clear(&self) {
93        self.entries.lock().clear();
94    }
95}
96
97impl Default for RotationCache {
98    fn default() -> Self {
99        Self::default_sized()
100    }
101}
102
103impl core::fmt::Debug for RotationCache {
104    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
105        // `entries` is intentionally reduced to its length — printing the
106        // full `RotationMatrix` payload would dump megabytes of f64 data.
107        f.debug_struct("RotationCache")
108            .field("capacity", &self.capacity)
109            .field("entries", &self.len())
110            .finish()
111    }
112}
tinyquant_core/codec/rotation_cache.rs

tinyquant_core/codec/
rotation_cache.rs
