1use std::collections::{HashMap, VecDeque};
7use std::hash::Hash;
8use std::num::NonZeroUsize;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
14pub enum EvictionPolicy {
15 #[default]
20 Lru,
21
22 Lfu,
27
28 Fifo,
32}
33
34pub(crate) trait EvictionStore<K, V>: Send {
36 fn get(&mut self, key: &K) -> Option<&V>;
38
39 fn insert(&mut self, key: K, value: V) -> Option<(K, V)>;
42
43 fn remove(&mut self, key: &K) -> Option<V>;
45
46 fn len(&self) -> usize;
48
49 fn clear(&mut self);
51
52 #[allow(dead_code)]
54 fn is_empty(&self) -> bool {
55 self.len() == 0
56 }
57}
58
59pub(crate) struct LruStore<K, V> {
61 cache: lru::LruCache<K, V>,
62}
63
64impl<K: Hash + Eq, V> LruStore<K, V> {
65 pub(crate) fn new(capacity: usize) -> Self {
66 let cap = NonZeroUsize::new(capacity).unwrap_or(NonZeroUsize::new(100).unwrap());
67 Self {
68 cache: lru::LruCache::new(cap),
69 }
70 }
71}
72
73impl<K: Hash + Eq + Send, V: Send> EvictionStore<K, V> for LruStore<K, V> {
74 fn get(&mut self, key: &K) -> Option<&V> {
75 self.cache.get(key)
76 }
77
78 fn insert(&mut self, key: K, value: V) -> Option<(K, V)> {
79 self.cache.push(key, value)
80 }
81
82 fn remove(&mut self, key: &K) -> Option<V> {
83 self.cache.pop(key)
84 }
85
86 fn len(&self) -> usize {
87 self.cache.len()
88 }
89
90 fn clear(&mut self) {
91 self.cache.clear();
92 }
93}
94
95pub(crate) struct LfuStore<K, V> {
97 data: HashMap<K, V>,
98 frequencies: HashMap<K, usize>,
99 capacity: usize,
100}
101
102impl<K: Hash + Eq + Clone, V> LfuStore<K, V> {
103 pub(crate) fn new(capacity: usize) -> Self {
104 Self {
105 data: HashMap::with_capacity(capacity),
106 frequencies: HashMap::with_capacity(capacity),
107 capacity: capacity.max(1),
108 }
109 }
110
111 fn find_lfu_key(&self) -> Option<K> {
112 self.frequencies
113 .iter()
114 .min_by_key(|(_, &freq)| freq)
115 .map(|(k, _)| k.clone())
116 }
117}
118
119impl<K: Hash + Eq + Clone + Send, V: Send> EvictionStore<K, V> for LfuStore<K, V> {
120 fn get(&mut self, key: &K) -> Option<&V> {
121 if self.data.contains_key(key) {
122 *self.frequencies.entry(key.clone()).or_insert(0) += 1;
123 self.data.get(key)
124 } else {
125 None
126 }
127 }
128
129 fn insert(&mut self, key: K, value: V) -> Option<(K, V)> {
130 if self.data.contains_key(&key) {
132 let old_value = self.data.insert(key.clone(), value)?;
133 *self.frequencies.entry(key.clone()).or_insert(0) += 1;
134 return Some((key, old_value));
135 }
136
137 let evicted = if self.data.len() >= self.capacity {
139 self.find_lfu_key().and_then(|lfu_key| {
140 let evicted_value = self.data.remove(&lfu_key)?;
141 self.frequencies.remove(&lfu_key);
142 Some((lfu_key, evicted_value))
143 })
144 } else {
145 None
146 };
147
148 self.data.insert(key.clone(), value);
150 self.frequencies.insert(key, 1);
151
152 evicted
153 }
154
155 fn remove(&mut self, key: &K) -> Option<V> {
156 self.frequencies.remove(key);
157 self.data.remove(key)
158 }
159
160 fn len(&self) -> usize {
161 self.data.len()
162 }
163
164 fn clear(&mut self) {
165 self.data.clear();
166 self.frequencies.clear();
167 }
168}
169
170pub(crate) struct FifoStore<K, V> {
172 data: HashMap<K, V>,
173 order: VecDeque<K>,
174 capacity: usize,
175}
176
177impl<K: Hash + Eq + Clone, V> FifoStore<K, V> {
178 pub(crate) fn new(capacity: usize) -> Self {
179 Self {
180 data: HashMap::with_capacity(capacity),
181 order: VecDeque::with_capacity(capacity),
182 capacity: capacity.max(1),
183 }
184 }
185}
186
187impl<K: Hash + Eq + Clone + Send, V: Send> EvictionStore<K, V> for FifoStore<K, V> {
188 fn get(&mut self, key: &K) -> Option<&V> {
189 self.data.get(key)
190 }
191
192 fn insert(&mut self, key: K, value: V) -> Option<(K, V)> {
193 if self.data.contains_key(&key) {
195 let old_value = self.data.insert(key.clone(), value)?;
196 return Some((key, old_value));
197 }
198
199 let evicted = if self.data.len() >= self.capacity {
201 self.order.pop_front().and_then(|old_key| {
202 let evicted_value = self.data.remove(&old_key)?;
203 Some((old_key, evicted_value))
204 })
205 } else {
206 None
207 };
208
209 self.data.insert(key.clone(), value);
211 self.order.push_back(key);
212
213 evicted
214 }
215
216 fn remove(&mut self, key: &K) -> Option<V> {
217 self.order.retain(|k| k != key);
218 self.data.remove(key)
219 }
220
221 fn len(&self) -> usize {
222 self.data.len()
223 }
224
225 fn clear(&mut self) {
226 self.data.clear();
227 self.order.clear();
228 }
229}
230
231#[cfg(test)]
232mod tests {
233 use super::*;
234
235 #[test]
236 fn test_lru_eviction() {
237 let mut store = LruStore::new(2);
238
239 store.insert("a", 1);
240 store.insert("b", 2);
241
242 assert_eq!(store.get(&"a"), Some(&1));
244
245 let evicted = store.insert("c", 3);
247 assert_eq!(evicted, Some(("b", 2)));
248
249 assert_eq!(store.get(&"a"), Some(&1));
250 assert_eq!(store.get(&"b"), None);
251 assert_eq!(store.get(&"c"), Some(&3));
252 }
253
254 #[test]
255 fn test_lfu_eviction() {
256 let mut store = LfuStore::new(2);
257
258 store.insert("a", 1);
259 store.insert("b", 2);
260
261 store.get(&"a");
263 store.get(&"a");
264 store.get(&"a");
265
266 store.get(&"b");
268
269 let evicted = store.insert("c", 3);
271 assert_eq!(evicted.map(|(k, _)| k), Some("b"));
272
273 assert_eq!(store.get(&"a"), Some(&1));
274 assert_eq!(store.get(&"b"), None);
275 assert_eq!(store.get(&"c"), Some(&3));
276 }
277
278 #[test]
279 fn test_fifo_eviction() {
280 let mut store = FifoStore::new(2);
281
282 store.insert("a", 1);
283 store.insert("b", 2);
284
285 store.get(&"b");
287 store.get(&"b");
288
289 let evicted = store.insert("c", 3);
291 assert_eq!(evicted, Some(("a", 1)));
292
293 assert_eq!(store.get(&"a"), None);
294 assert_eq!(store.get(&"b"), Some(&2));
295 assert_eq!(store.get(&"c"), Some(&3));
296 }
297
298 #[test]
299 fn test_eviction_policy_default() {
300 assert_eq!(EvictionPolicy::default(), EvictionPolicy::Lru);
301 }
302}