tiny_cache/
lru_cache.rs

1use std::{
2    hash::{
3        Hash,
4        Hasher,
5    },
6    rc::Rc,
7};
8
9use rustc_hash::{
10    FxHashMap,
11    FxHasher,
12};
13use smallvec::SmallVec;
14
15/// Reuse costly results across users, such as text measurements shared by several labels.
16/// Each user keeps its last `USER_CAPACITY` results. Saved results leave when no user needs them.
17pub struct LRUCache<V, ID: Hash, const USER_CAPACITY: usize = 2> {
18    map: FxHashMap<u64, (i32, Rc<V>)>,
19    /// Tracks which saved results each user still needs.
20    pub users: FxHashMap<ID, SmallVec<[u64; 2]>>,
21}
22
23impl<V, ID: Hash, const USER_CAPACITY: usize> Default for LRUCache<V, ID, USER_CAPACITY> {
24    fn default() -> Self {
25        Self {
26            map: FxHashMap::default(),
27            users: FxHashMap::default(),
28        }
29    }
30}
31
32impl<V, ID: Hash + Eq, const USER_CAPACITY: usize> LRUCache<V, ID, USER_CAPACITY> {
33    fn hash_of<H: Hash>(hash_value: &H) -> u64 {
34        let mut hasher = FxHasher::default();
35        hash_value.hash(&mut hasher);
36        hasher.finish()
37    }
38
39    /// Reuses a saved result for this user and makes room by letting go of their oldest results.
40    /// If it returns `None`, create the result and save it with [`insert`](Self::insert).
41    pub fn utilize<H: Hash>(&mut self, id: ID, hash_value: &H) -> Option<Rc<V>> {
42        let hash = Self::hash_of(hash_value);
43        let entry = self.map.get_mut(&hash);
44        let cached_value = entry.as_ref().map(|(_, value)| value.clone());
45
46        let hashes = self.users.entry(id).or_default();
47
48        if let Some(position) = hashes.iter().position(|used| *used == hash) {
49            hashes.remove(position);
50            hashes.push(hash);
51
52            return cached_value;
53        }
54
55        if let Some((users, _)) = entry {
56            *users += 1;
57            hashes.push(hash);
58        }
59
60        let keep = if cached_value.is_some() {
61            USER_CAPACITY
62        } else {
63            USER_CAPACITY.saturating_sub(1)
64        };
65        let extra = hashes.len().saturating_sub(keep);
66
67        for old_hash in hashes.drain(0..extra) {
68            let Some((users, _)) = self.map.get_mut(&old_hash) else {
69                continue;
70            };
71
72            *users -= 1;
73
74            if *users == 0 {
75                self.map.remove(&old_hash);
76            }
77        }
78
79        cached_value
80    }
81
82    /// Saves a newly created result for this user so others can reuse it without repeating the work.
83    /// Call [`utilize`](Self::utilize) first to check for a saved result and make room for this one.
84    pub fn insert<H: Hash>(&mut self, id: ID, hash_value: &H, value: V) -> Rc<V> {
85        let hash = Self::hash_of(hash_value);
86        let value = Rc::new(value);
87
88        let (users, _) = self.map.entry(hash).or_insert_with(|| (0, value.clone()));
89        let user_hashes = self.users.entry(id).or_default();
90
91        if !user_hashes.contains(&hash) {
92            user_hashes.push(hash);
93            *users += 1;
94        }
95
96        value
97    }
98
99    /// Counts saved results, counting a result shared by several users only once.
100    pub fn len(&self) -> usize {
101        self.map.len()
102    }
103
104    pub fn is_empty(&self) -> bool {
105        self.map.is_empty()
106    }
107
108    /// Counts users being tracked, including those with no saved results yet.
109    pub fn users_len(&self) -> usize {
110        self.users.len()
111    }
112
113    /// Stops keeping results for a user that has gone away, such as a label removed from the screen.
114    /// Results still needed by other users stay saved, and copies already handed out remain usable.
115    pub fn remove(&mut self, id: &ID) {
116        let Some(hashes) = self.users.remove(id) else {
117            return;
118        };
119
120        for hash in hashes {
121            let Some((users, _)) = self.map.get_mut(&hash) else {
122                continue;
123            };
124
125            *users -= 1;
126
127            if *users == 0 {
128                self.map.remove(&hash);
129            }
130        }
131    }
132
133    /// Clears saved results and user tracking.
134    pub fn reset(&mut self) {
135        self.map.clear();
136        self.users.clear();
137    }
138}
139
140impl<V, ID: Hash, const USER_CAPACITY: usize> std::fmt::Debug for LRUCache<V, ID, USER_CAPACITY> {
141    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
142        f.debug_struct("LRUCache")
143            .field("cached_values", &self.map.len())
144            .field("cached_users", &self.users.len())
145            .finish()
146    }
147}
148
149#[cfg(test)]
150mod test {
151    use std::rc::Rc;
152
153    use crate::lru_cache::LRUCache;
154
155    type Cache = LRUCache<i32, u64>;
156
157    #[test]
158    fn lru_cache() {
159        let mut cache = Cache::default();
160
161        cache
162            .utilize(1, &50)
163            .unwrap_or_else(|| cache.insert(1, &50, 500));
164        assert_eq!(cache.utilize(1, &50), Some(Rc::new(500)));
165
166        cache
167            .utilize(1, &60)
168            .unwrap_or_else(|| cache.insert(1, &60, 600));
169        assert_eq!(cache.utilize(1, &60), Some(Rc::new(600)));
170        assert_eq!(cache.utilize(1, &50), Some(Rc::new(500)));
171
172        cache
173            .utilize(1, &70)
174            .unwrap_or_else(|| cache.insert(1, &70, 700));
175        assert!(!cache.map.contains_key(&Cache::hash_of(&60)));
176        assert!(cache.map.contains_key(&Cache::hash_of(&50)));
177
178        cache.remove(&1);
179        assert!(!cache.map.contains_key(&Cache::hash_of(&50)));
180        assert!(!cache.map.contains_key(&Cache::hash_of(&70)));
181
182        cache
183            .utilize(1, &70)
184            .unwrap_or_else(|| cache.insert(1, &70, 700));
185        assert!(cache.utilize(2, &70).is_some());
186
187        cache.remove(&1);
188        assert!(cache.map.contains_key(&Cache::hash_of(&70)));
189        cache.remove(&2);
190        assert!(!cache.map.contains_key(&Cache::hash_of(&70)));
191    }
192
193    #[test]
194    fn keeps_both_hashes_when_another_user_had_one() {
195        let mut cache = Cache::default();
196
197        cache
198            .utilize(1, &10)
199            .unwrap_or_else(|| cache.insert(1, &10, 100));
200        cache
201            .utilize(2, &20)
202            .unwrap_or_else(|| cache.insert(2, &20, 200));
203        cache
204            .utilize(1, &20)
205            .unwrap_or_else(|| cache.insert(1, &20, 200));
206
207        assert_eq!(cache.users[&1].len(), 2);
208        assert!(cache.map.contains_key(&Cache::hash_of(&10)));
209        assert!(cache.map.contains_key(&Cache::hash_of(&20)));
210    }
211
212    #[test]
213    fn only_keeps_two_hashes_per_user() {
214        let mut cache = Cache::default();
215
216        for hash in [10, 20, 30, 40] {
217            let id = 1000 + hash as u64;
218            cache
219                .utilize(id, &hash)
220                .unwrap_or_else(|| cache.insert(id, &hash, hash * 10));
221        }
222
223        for hash in [10, 20, 30, 40] {
224            cache
225                .utilize(1, &hash)
226                .unwrap_or_else(|| cache.insert(1, &hash, hash * 10));
227        }
228
229        assert_eq!(cache.users[&1].len(), 2);
230    }
231
232    #[test]
233    fn using_a_hash_again_keeps_it() {
234        let mut cache = Cache::default();
235
236        for hash in [10, 20, 10, 30] {
237            cache
238                .utilize(1, &hash)
239                .unwrap_or_else(|| cache.insert(1, &hash, hash * 10));
240        }
241
242        assert!(cache.map.contains_key(&Cache::hash_of(&10)));
243        assert!(!cache.map.contains_key(&Cache::hash_of(&20)));
244        assert!(cache.map.contains_key(&Cache::hash_of(&30)));
245    }
246}
tiny_cache/lru_cache.rs

tiny_cache/
lru_cache.rs
