ps_rclru/
lib.rs

1use std::{
2    collections::{HashMap, LinkedList},
3    hash::Hash,
4    rc::Rc,
5};
6
7pub struct LRU<K, V> {
8    list: LinkedList<(Rc<K>, Rc<V>)>,
9    map: HashMap<Rc<K>, (Rc<V>, usize)>,
10    num_items: usize,
11    max_items: usize,
12}
13
14impl<K: Eq + Hash, V> LRU<K, V> {
15    pub fn new(max_items: usize) -> Self {
16        Self {
17            list: LinkedList::new(),
18            map: HashMap::new(),
19            num_items: 0,
20            max_items,
21        }
22    }
23
24    /// **Each key must map to exactly one value.** Pushing multiple different value for the same key is undefined behaviour.
25    pub fn push(&mut self, key: Rc<K>, value: Rc<V>) -> Option<(Rc<K>, Rc<V>)> {
26        if let Some((_, count)) = self.map.get_mut(&key) {
27            // value already in LRU
28            self.list.push_back((key, value));
29            *count += 1;
30
31            return self.maybe_gc();
32        }
33
34        // new element inserted
35        self.list.push_back((key.clone(), value.clone()));
36        self.map.insert(key, (value, 1));
37        self.num_items += 1;
38
39        return self.maybe_gc();
40    }
41
42    #[inline(always)]
43    pub fn maybe_gc(&mut self) -> Option<(Rc<K>, Rc<V>)> {
44        if self.num_items > self.max_items {
45            self.gc()
46        } else {
47            None
48        }
49    }
50
51    pub fn gc(&mut self) -> Option<(Rc<K>, Rc<V>)> {
52        let mut iterations = 0;
53
54        while iterations < self.list.len() && self.num_items > self.max_items {
55            iterations += 1;
56
57            if let Some((key, value)) = self.list.pop_front() {
58                if let Some((_, count)) = self.map.get_mut(&key) {
59                    if *count > 1 {
60                        // multiple references exist in list
61                        *count -= 1;
62                    } else if Rc::strong_count(&value) > 2 {
63                        // a reference exists outside this LRU cache
64                        self.list.push_back((key, value));
65                    } else {
66                        // evict from cache
67                        self.map.remove(&key);
68                        self.num_items -= 1;
69                        // return the evicted pair
70                        return Some((key, value));
71                    }
72                } else {
73                    return Some((key, value));
74                }
75            } else {
76                // list is empty
77                return None;
78            }
79        }
80
81        return None;
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88    use std::rc::Rc;
89
90    #[test]
91    fn test_basic_insertion_and_eviction() {
92        let mut lru = LRU {
93            list: LinkedList::new(),
94            map: HashMap::new(),
95            num_items: 0,
96            max_items: 3,
97        };
98
99        // Insert elements
100        assert!(lru.push(Rc::new(1), Rc::new("a")).is_none());
101        assert!(lru.push(Rc::new(2), Rc::new("b")).is_none());
102        assert!(lru.push(Rc::new(3), Rc::new("c")).is_none());
103
104        // Check no eviction yet
105        assert_eq!(lru.num_items, 3);
106
107        // Insert one more element, triggering eviction
108        let evicted = lru.push(Rc::new(4), Rc::new("d")).unwrap();
109        assert_eq!(*evicted.0, 1); // LRU element "1" should be evicted
110
111        // Check internal state
112        assert_eq!(lru.num_items, 3);
113        assert!(lru.map.contains_key(&Rc::new(2)));
114        assert!(lru.map.contains_key(&Rc::new(3)));
115        assert!(lru.map.contains_key(&Rc::new(4)));
116    }
117
118    #[test]
119    fn test_reinsertion_of_existing_key() {
120        let mut lru = LRU {
121            list: LinkedList::new(),
122            map: HashMap::new(),
123            num_items: 0,
124            max_items: 3,
125        };
126
127        let key = Rc::new(1);
128        let value = Rc::new("a");
129
130        lru.push(key.clone(), value.clone());
131        lru.push(Rc::new(2), Rc::new("b"));
132        lru.push(Rc::new(3), Rc::new("c"));
133
134        // Reinserting existing key should update its position
135        lru.push(key.clone(), value.clone());
136
137        // Insert another element to trigger eviction
138        let evicted = lru.push(Rc::new(4), Rc::new("d")).unwrap();
139        assert_eq!(*evicted.0, 2); // Element "2" should be evicted as "1" was recently used
140
141        // Check internal state
142        assert_eq!(lru.num_items, 3);
143        assert!(lru.map.contains_key(&key));
144        assert!(lru.map.contains_key(&Rc::new(3)));
145        assert!(lru.map.contains_key(&Rc::new(4)));
146    }
147
148    #[test]
149    fn test_gc_with_reference_counts() {
150        let mut lru = LRU {
151            list: LinkedList::new(),
152            map: HashMap::new(),
153            num_items: 0,
154            max_items: 2,
155        };
156
157        let key1 = Rc::new(1);
158        let value1 = Rc::new("a");
159        lru.push(key1.clone(), value1.clone());
160
161        let key2 = Rc::new(2);
162        let value2 = Rc::new("b");
163        lru.push(key2.clone(), value2.clone());
164
165        // External reference to key1 and value1
166        let _key1_ref = Rc::clone(&key1);
167        let _value1_ref = Rc::clone(&value1);
168
169        let key3 = Rc::new(3);
170        let value3 = Rc::new("c");
171
172        // Adding new element should trigger GC
173        let evicted = lru.push(key3.clone(), value3.clone());
174        assert!(evicted.is_none()); // No eviction as key1/value1 has external refs
175
176        // Drop external reference to value2
177        drop(value2);
178
179        // Add another element, this should evict key2/value2
180        let evicted = lru.push(Rc::new(4), Rc::new("d")).unwrap();
181        assert_eq!(*evicted.0, 2); // Key "2" should be evicted
182
183        // Check internal state
184        assert_eq!(lru.num_items, 3);
185        assert!(lru.map.contains_key(&Rc::new(1)));
186        assert!(lru.map.contains_key(&Rc::new(3)));
187        assert!(lru.map.contains_key(&Rc::new(4)));
188    }
189
190    #[test]
191    fn test_handling_empty_lru() {
192        let mut lru = LRU {
193            list: LinkedList::new(),
194            map: HashMap::new(),
195            num_items: 0,
196            max_items: 2,
197        };
198
199        // Try GC on empty LRU
200        let evicted = lru.gc();
201        assert!(evicted.is_none());
202
203        // Insert and remove an element
204        lru.push(Rc::new(1), Rc::new("a"));
205        let evicted = lru.push(Rc::new(2), Rc::new("b"));
206        assert!(evicted.is_none());
207
208        // Remove all elements
209        let _ = lru.push(Rc::new(3), Rc::new("c"));
210        let _ = lru.push(Rc::new(4), Rc::new("d"));
211        assert_eq!(lru.num_items, 2);
212    }
213
214    #[test]
215    fn test_push_new_item() {
216        let mut lru = LRU::new(2);
217        let key = Rc::new(1);
218        let value = Rc::new("a");
219
220        let evicted = lru.push(key.clone(), value.clone());
221        assert!(evicted.is_none());
222        assert_eq!(lru.num_items, 1);
223        assert!(lru.map.contains_key(&key));
224    }
225
226    #[test]
227    fn test_eviction() {
228        let mut lru = LRU::new(2);
229        let key1 = Rc::new(1);
230        let key2 = Rc::new(2);
231        let key3 = Rc::new(3);
232        let value1 = Rc::new("a");
233        let value2 = Rc::new("b");
234        let value3 = Rc::new("c");
235
236        lru.push(key1, value1);
237        lru.push(key2.clone(), value2.clone());
238        let evicted = lru.push(key3.clone(), value3.clone());
239
240        assert!(evicted.is_some());
241        let (evicted_key, evicted_value) = evicted.unwrap();
242        assert_eq!(*evicted_key, 1);
243        assert_eq!(*evicted_value, "a");
244        assert_eq!(lru.num_items, 2);
245        assert!(lru.map.contains_key(&key2));
246        assert!(lru.map.contains_key(&key3));
247    }
248
249    #[test]
250    fn test_reference_count_handling() {
251        let mut lru = LRU::new(1);
252        let key: Rc<i32> = Rc::new(1);
253        let value = Rc::new("a");
254
255        {
256            let key_ref = key.clone();
257            let value_ref = value.clone();
258            lru.push(key_ref, value_ref);
259            assert_eq!(Rc::strong_count(&key), 3);
260            assert_eq!(Rc::strong_count(&value), 3);
261        }
262
263        let evicted = lru.push(key.clone(), value.clone());
264        assert!(evicted.is_none());
265        assert_eq!(Rc::strong_count(&key), 4);
266        assert_eq!(Rc::strong_count(&value), 4);
267    }
268
269    #[test]
270    fn test_gc_on_empty_list() {
271        let mut lru = LRU::<i32, Box<dyn std::any::Any>>::new(1);
272        let evicted = lru.gc();
273        assert!(evicted.is_none());
274    }
275
276    fn setup_lru(max_items: usize) -> LRU<i32, i32> {
277        LRU {
278            list: LinkedList::new(),
279            map: HashMap::new(),
280            num_items: 0,
281            max_items,
282        }
283    }
284
285    #[test]
286    fn test_push_new_element() {
287        let mut lru = setup_lru(2);
288        let key = Rc::new(1);
289        let value = Rc::new(10);
290
291        assert_eq!(lru.push(Rc::clone(&key), Rc::clone(&value)), None);
292        assert_eq!(lru.list.len(), 1);
293        assert_eq!(lru.map.len(), 1);
294        assert_eq!(lru.num_items, 1);
295    }
296
297    #[test]
298    fn test_push_existing_element() {
299        let mut lru = setup_lru(2);
300        let key = Rc::new(1);
301        let value = Rc::new(10);
302
303        lru.push(Rc::clone(&key), Rc::clone(&value));
304        assert_eq!(lru.push(Rc::clone(&key), Rc::clone(&value)), None);
305        assert_eq!(lru.list.len(), 2);
306        assert_eq!(lru.map.len(), 1);
307        assert_eq!(lru.num_items, 1);
308    }
309
310    #[test]
311    fn test_gc_multiple_references() {
312        let mut lru = setup_lru(2);
313        let key1 = Rc::new(1);
314        let value1 = Rc::new(10);
315        let key2 = Rc::new(2);
316        let value2 = Rc::new(20);
317
318        lru.push(Rc::clone(&key1), Rc::clone(&value1));
319        lru.push(Rc::clone(&key2), Rc::clone(&value2));
320        lru.push(Rc::clone(&key1), Rc::clone(&value1)); // key1 is pushed again
321
322        assert_eq!(lru.gc(), None); // No eviction should happen
323        assert_eq!(lru.list.len(), 3);
324        assert_eq!(lru.map.len(), 2);
325        assert_eq!(lru.num_items, 2);
326    }
327
328    #[test]
329    fn test_push_and_retrieve() {
330        let mut lru = LRU {
331            list: LinkedList::new(),
332            map: HashMap::new(),
333            num_items: 0,
334            max_items: 3,
335        };
336
337        let key1 = Rc::new(1);
338        let value1 = Rc::new("a");
339        let key2 = Rc::new(2);
340        let value2 = Rc::new("b");
341        let key3 = Rc::new(3);
342        let value3 = Rc::new("c");
343
344        lru.push(key1.clone(), value1.clone());
345        lru.push(key2.clone(), value2.clone());
346        lru.push(key3.clone(), value3.clone());
347
348        assert_eq!(lru.num_items, 3);
349        assert!(lru.map.contains_key(&key1));
350        assert!(lru.map.contains_key(&key2));
351        assert!(lru.map.contains_key(&key3));
352    }
353
354    #[test]
355    fn test_eviction_policy() {
356        let mut lru = LRU {
357            list: LinkedList::new(),
358            map: HashMap::new(),
359            num_items: 0,
360            max_items: 2,
361        };
362
363        let key1 = Rc::new(1);
364        let value1 = Rc::new("a");
365        let key2 = Rc::new(2);
366        let value2 = Rc::new("b");
367        let key3 = Rc::new(3);
368        let value3 = Rc::new("c");
369
370        lru.push(key1.clone(), value1.clone());
371        lru.push(key2.clone(), value2.clone());
372        let evicted = lru.push(key3.clone(), value3.clone());
373
374        assert_eq!(lru.num_items, 3);
375        assert!(lru.map.contains_key(&key1));
376        assert!(lru.map.contains_key(&key2));
377        assert!(lru.map.contains_key(&key3));
378        assert_eq!(evicted, None);
379    }
380
381    #[test]
382    fn test_multiple_references() {
383        let mut lru = LRU {
384            list: LinkedList::new(),
385            map: HashMap::new(),
386            num_items: 0,
387            max_items: 2,
388        };
389
390        let key1 = Rc::new(1);
391        let value1 = Rc::new("a");
392
393        lru.push(key1.clone(), value1.clone());
394        lru.push(key1.clone(), value1.clone());
395
396        assert_eq!(lru.num_items, 1);
397        assert_eq!(lru.list.len(), 2);
398        assert_eq!(lru.map.get(&key1).unwrap().1, 2);
399    }
400
401    #[test]
402    fn test_eviction_with_external_references() {
403        let mut lru = LRU {
404            list: LinkedList::new(),
405            map: HashMap::new(),
406            num_items: 0,
407            max_items: 1,
408        };
409
410        let key1 = Rc::new(1);
411        let value1 = Rc::new("a");
412        let key2 = Rc::new(2);
413        let value2 = Rc::new("b");
414
415        lru.push(key1, value1);
416        let evicted = lru.push(key2.clone(), value2.clone());
417
418        assert_eq!(lru.num_items, 1);
419        assert!(lru.map.contains_key(&key2));
420        assert_eq!(evicted, Some((Rc::new(1), Rc::new("a"))));
421    }
422
423    #[test]
424    fn test_gc_behavior() {
425        let mut lru = LRU {
426            list: LinkedList::new(),
427            map: HashMap::new(),
428            num_items: 0,
429            max_items: 1,
430        };
431
432        let key1 = Rc::new(1);
433        let value1 = Rc::new("a");
434        let key2 = Rc::new(2);
435        let value2 = Rc::new("b");
436
437        lru.push(key1.clone(), value1.clone());
438        lru.push(key2.clone(), value2.clone());
439
440        let gc_result = lru.gc();
441
442        assert!(gc_result.is_none());
443        assert_eq!(lru.num_items, 2);
444        assert!(lru.map.contains_key(&key1));
445        assert!(lru.map.contains_key(&key2));
446    }
447
448    #[test]
449    fn test_push_within_limit() {
450        let mut lru = LRU::new(3);
451        let k1 = Rc::new(1);
452        let v1 = Rc::new(10);
453        assert_eq!(lru.push(Rc::clone(&k1), Rc::clone(&v1)), None);
454        assert_eq!(lru.num_items, 1);
455    }
456
457    #[test]
458    fn test_push_eviction() {
459        let mut lru = LRU::new(2);
460        let k1 = Rc::new(1);
461        let v1 = Rc::new(10);
462        let k2 = Rc::new(2);
463        let v2 = Rc::new(20);
464        let k3 = Rc::new(3);
465        let v3 = Rc::new(30);
466
467        lru.push(k1, v1);
468        lru.push(Rc::clone(&k2), Rc::clone(&v2));
469        let evicted = lru.push(Rc::clone(&k3), Rc::clone(&v3));
470        assert!(evicted.is_some());
471        assert_eq!(*evicted.unwrap().0, 1);
472        assert_eq!(lru.num_items, 2);
473    }
474
475    #[test]
476    fn test_gc_no_eviction() {
477        let mut lru = LRU::new(3);
478        let k1 = Rc::new(1);
479        let v1 = Rc::new(10);
480        lru.push(Rc::clone(&k1), Rc::clone(&v1));
481        assert_eq!(lru.gc(), None);
482    }
483
484    #[test]
485    fn test_gc_with_eviction() {
486        let mut lru = LRU::new(1);
487        let k1 = Rc::new(1);
488        let v1 = Rc::new(10);
489        let k2 = Rc::new(2);
490        let v2 = Rc::new(20);
491
492        lru.push(Rc::clone(&k1), Rc::clone(&v1));
493        lru.push(Rc::clone(&k2), Rc::clone(&v2));
494
495        drop(v1);
496
497        let evicted = lru.gc();
498
499        assert!(evicted.is_some());
500        assert_eq!(*evicted.unwrap().0, 1);
501    }
502}
ps_rclru/lib.rs

ps_rclru/
lib.rs
