near_cache/
cell.rs

1use lru::LruCache;
2use std::borrow::Borrow;
3use std::cell::RefCell;
4use std::convert::Infallible;
5use std::hash::Hash;
6use std::num::NonZeroUsize;
7
8/// A wrapper around `LruCache` to provide shared `&` access to content.
9pub struct CellLruCache<K, V> {
10    inner: RefCell<LruCache<K, V>>,
11}
12
13impl<K, V> CellLruCache<K, V>
14where
15    K: Hash + Eq,
16    V: Clone,
17{
18    /// Creates a new `LRU` cache that holds at most `cap` items.
19    pub fn new(cap: usize) -> Self {
20        Self { inner: RefCell::new(LruCache::<K, V>::new(NonZeroUsize::new(cap).unwrap())) }
21    }
22
23    /// Returns the number of key-value pairs that are currently in the cache.
24    pub fn len(&self) -> usize {
25        self.inner.borrow().len()
26    }
27
28    /// Returns true if the cache is empty and false otherwise
29    pub fn is_empty(&self) -> bool {
30        self.inner.borrow().is_empty()
31    }
32
33    /// Return the value of they key in the cache otherwise computes the value and inserts it into
34    /// the cache. If the key is already in the cache, they get moved to the head of
35    /// the LRU list.
36    pub fn get_or_put<F>(&self, key: K, f: F) -> V
37    where
38        V: Clone,
39        F: FnOnce(&K) -> V,
40    {
41        Result::<_, Infallible>::unwrap(self.get_or_try_put(key, |k| Ok(f(k))))
42    }
43
44    /// Returns the value of they key in the cache if present, otherwise
45    /// computes the value using the provided closure.
46    ///
47    /// If the key is already in the cache, it gets moved to the head of the LRU
48    /// list.
49    ///
50    /// If the provided closure fails, the error is returned and the cache is
51    /// not updated.
52    pub fn get_or_try_put<F, E>(&self, key: K, f: F) -> Result<V, E>
53    where
54        V: Clone,
55        F: FnOnce(&K) -> Result<V, E>,
56    {
57        if let Some(result) = self.get(&key) {
58            return Ok(result);
59        }
60        let val = f(&key)?;
61        let val_clone = val.clone();
62        self.inner.borrow_mut().put(key, val_clone);
63        Ok(val)
64    }
65
66    /// Puts a key-value pair into cache. If the key already exists in the cache,
67    /// then it updates the key's value.
68    pub fn put(&self, key: K, value: V) {
69        self.inner.borrow_mut().put(key, value);
70    }
71
72    pub fn pop<Q>(&self, key: &Q) -> Option<V>
73    where
74        K: Borrow<Q>,
75        Q: Hash + Eq + ?Sized,
76    {
77        self.inner.borrow_mut().pop(key)
78    }
79
80    /// Returns the value of the key in the cache or None if it is not present in the cache.
81    /// Moves the key to the head of the LRU list if it exists.
82    pub fn get<Q>(&self, key: &Q) -> Option<V>
83    where
84        K: Borrow<Q>,
85        Q: Hash + Eq + ?Sized,
86    {
87        self.inner.borrow_mut().get(key).cloned()
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn test_cache() {
97        let cache = CellLruCache::<u64, Vec<u64>>::new(100);
98
99        assert_eq!(cache.get(&0u64), None);
100        assert_eq!(cache.get_or_put(123u64, |key| vec![*key, 123]), vec![123u64, 123]);
101        assert_eq!(cache.get(&123u64), Some(vec![123u64, 123]));
102        assert_eq!(cache.get(&0u64), None);
103    }
104}