freya_core/
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
15pub struct LRUCache<V, ID: Hash> {
16    map: FxHashMap<u64, (i32, Rc<V>)>,
17    pub users: FxHashMap<ID, SmallVec<[u64; 2]>>,
18}
19
20impl<V, ID: Hash> Default for LRUCache<V, ID> {
21    fn default() -> Self {
22        Self {
23            map: FxHashMap::default(),
24            users: FxHashMap::default(),
25        }
26    }
27}
28
29impl<V, ID: Hash + Eq> LRUCache<V, ID> {
30    pub fn get<H: Hash>(&self, hash_value: &H) -> Option<Rc<V>> {
31        let mut hasher = FxHasher::default();
32        hash_value.hash(&mut hasher);
33        let hash = hasher.finish();
34        let value = self.map.get(&hash);
35
36        value.as_ref().map(|v| v.1.clone())
37    }
38
39    pub fn utilize<H: Hash>(&mut self, id: ID, hash_value: &H) -> Option<Rc<V>> {
40        let mut hasher = FxHasher::default();
41        hash_value.hash(&mut hasher);
42        let hash = hasher.finish();
43        let mut value = self.map.get_mut(&hash);
44
45        let cache_value = value.as_ref().map(|v| v.1.clone());
46
47        let hashes = self.users.entry(id).or_default();
48
49        // New hashed value
50        if !hashes.contains(&hash) {
51            if let Some(value) = &mut value {
52                value.0 += 1;
53                hashes.push(hash);
54            }
55
56            let index = match hashes.len() {
57                ..=1 => 0,
58                2 => 1,
59                len => len - 2,
60            };
61            // Clean the first current hash
62            for old_hash in hashes.drain(0..index) {
63                let Some(entry) = self.map.get_mut(&old_hash) else {
64                    continue;
65                };
66
67                entry.0 -= 1;
68
69                if entry.0 == 0 {
70                    self.map.remove(&old_hash);
71                }
72            }
73        }
74
75        cache_value
76    }
77
78    pub fn insert<H: Hash>(&mut self, id: ID, hash_value: &H, value: V) -> Rc<V> {
79        let mut hasher = FxHasher::default();
80        hash_value.hash(&mut hasher);
81        let hash = hasher.finish();
82        let value = Rc::new(value);
83
84        self.map.entry(hash).or_insert_with(|| (0, value.clone()));
85
86        let user_hashes = self.users.entry(id).or_default();
87        if !user_hashes.contains(&hash) {
88            user_hashes.push(hash);
89            self.map.get_mut(&hash).unwrap().0 += 1;
90        }
91
92        value
93    }
94
95    pub fn remove(&mut self, id: &ID) {
96        let Some(hashes) = self.users.remove(id) else {
97            return;
98        };
99
100        for hash in hashes.iter() {
101            let Some(entry) = self.map.get_mut(hash) else {
102                continue;
103            };
104
105            entry.0 -= 1;
106
107            if entry.0 == 0 {
108                self.map.remove(hash);
109            }
110        }
111    }
112
113    pub fn reset(&mut self) {
114        self.map.clear();
115        self.users.clear();
116    }
117
118    pub fn print_metrics(&self) {
119        println!("Cached Values {}", self.map.len());
120        println!("Cache Users {}", self.users.len());
121    }
122}
123
124#[cfg(test)]
125mod test {
126    use std::rc::Rc;
127
128    use crate::lru_cache::LRUCache;
129
130    #[test]
131    fn lru_cache() {
132        let mut cache = LRUCache::<i32, u64>::default();
133
134        cache
135            .utilize(1, &50)
136            .unwrap_or_else(|| cache.insert(1, &50, 5000));
137
138        assert_eq!(cache.utilize(1, &50), Some(Rc::new(5000)));
139
140        cache
141            .utilize(1, &60)
142            .unwrap_or_else(|| cache.insert(1, &60, 6000));
143        assert_eq!(cache.utilize(1, &60), Some(Rc::new(6000)));
144        assert_eq!(cache.utilize(1, &50), Some(Rc::new(5000)));
145
146        cache
147            .utilize(1, &70)
148            .unwrap_or_else(|| cache.insert(1, &70, 7000));
149        assert!(cache.get(&50).is_none());
150
151        assert_eq!(cache.utilize(1, &60), Some(Rc::new(6000)));
152        assert_eq!(cache.utilize(1, &70), Some(Rc::new(7000)));
153        assert_eq!(cache.utilize(1, &60), Some(Rc::new(6000)));
154        assert_eq!(cache.utilize(1, &70), Some(Rc::new(7000)));
155
156        cache.remove(&1);
157        assert!(cache.get(&60).is_none());
158        assert!(cache.get(&70).is_none());
159
160        cache
161            .utilize(1, &70)
162            .unwrap_or_else(|| cache.insert(1, &70, 7000));
163        assert!(cache.utilize(2, &70).is_some());
164
165        cache.remove(&1);
166        assert!(cache.get(&70).is_some());
167        cache.remove(&2);
168        assert!(cache.get(&70).is_none());
169    }
170}
freya_core/lru_cache.rs

freya_core/
lru_cache.rs
