1use alloc::collections::VecDeque;
2use core::borrow::Borrow;
3use core::hash::Hash;
4use std::collections::hash_map::Entry;
5use std::collections::HashMap;
6
7/// A HashMap-alike, which never gets larger than a specified
8/// capacity, and evicts the oldest insertion to maintain this.
9///
10/// The requested capacity may be rounded up by the underlying
11/// collections. This implementation uses all the allocated
12/// storage.
13///
14/// This is inefficient: it stores keys twice.
15pub(crate) struct LimitedCache<K: Clone + Hash + Eq, V> {
16 map: HashMap<K, V>,
17
18 // first item is the oldest key
19 oldest: VecDeque<K>,
20}
21
22impl<K, V> LimitedCache<K, V>
23where
24 K: Eq + Hash + Clone + core::fmt::Debug,
25 V: Default,
26{
27 /// Create a new LimitedCache with the given rough capacity.
28 pub(crate) fn new(capacity_order_of_magnitude: usize) -> Self {
29 Self {
30 map: HashMap::with_capacity(capacity_order_of_magnitude),
31 oldest: VecDeque::with_capacity(capacity_order_of_magnitude),
32 }
33 }
34
35 pub(crate) fn get_or_insert_default_and_edit(&mut self, k: K, edit: impl FnOnce(&mut V)) {
36 let inserted_new_item = match self.map.entry(k) {
37 Entry::Occupied(value) => {
38 edit(value.into_mut());
39 false
40 }
41 entry @ Entry::Vacant(_) => {
42 self.oldest
43 .push_back(entry.key().clone());
44 edit(entry.or_insert_with(V::default));
45 true
46 }
47 };
48
49 // ensure next insertion does not require a realloc
50 if inserted_new_item && self.oldest.capacity() == self.oldest.len() {
51 if let Some(oldest_key) = self.oldest.pop_front() {
52 self.map.remove(&oldest_key);
53 }
54 }
55 }
56
57 pub(crate) fn insert(&mut self, k: K, v: V) {
58 let inserted_new_item = match self.map.entry(k) {
59 Entry::Occupied(mut old) => {
60 // Note: does not freshen entry in `oldest`
61 old.insert(v);
62 false
63 }
64
65 entry @ Entry::Vacant(_) => {
66 self.oldest
67 .push_back(entry.key().clone());
68 entry.or_insert(v);
69 true
70 }
71 };
72
73 // ensure next insertion does not require a realloc
74 if inserted_new_item && self.oldest.capacity() == self.oldest.len() {
75 if let Some(oldest_key) = self.oldest.pop_front() {
76 self.map.remove(&oldest_key);
77 }
78 }
79 }
80
81 pub(crate) fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
82 where
83 K: Borrow<Q>,
84 Q: Hash + Eq,
85 {
86 self.map.get(k)
87 }
88
89 pub(crate) fn get_mut<Q: ?Sized>(&mut self, k: &Q) -> Option<&mut V>
90 where
91 K: Borrow<Q>,
92 Q: Hash + Eq,
93 {
94 self.map.get_mut(k)
95 }
96
97 pub(crate) fn remove<Q: ?Sized>(&mut self, k: &Q) -> Option<V>
98 where
99 K: Borrow<Q>,
100 Q: Hash + Eq,
101 {
102 if let Some(value) = self.map.remove(k) {
103 // O(N) search, followed by O(N) removal
104 if let Some(index) = self
105 .oldest
106 .iter()
107 .position(|item| item.borrow() == k)
108 {
109 self.oldest.remove(index);
110 }
111 Some(value)
112 } else {
113 None
114 }
115 }
116}
117
118#[cfg(test)]
119mod tests {
120 type Test = super::LimitedCache<String, usize>;
121
122 #[test]
123 fn test_updates_existing_item() {
124 let mut t = Test::new(3);
125 t.insert("abc".into(), 1);
126 t.insert("abc".into(), 2);
127 assert_eq!(t.get("abc"), Some(&2));
128 }
129
130 #[test]
131 fn test_evicts_oldest_item() {
132 let mut t = Test::new(3);
133 t.insert("abc".into(), 1);
134 t.insert("def".into(), 2);
135 t.insert("ghi".into(), 3);
136
137 assert_eq!(t.get("abc"), None);
138 assert_eq!(t.get("def"), Some(&2));
139 assert_eq!(t.get("ghi"), Some(&3));
140 }
141
142 #[test]
143 fn test_evicts_second_oldest_item_if_first_removed() {
144 let mut t = Test::new(3);
145 t.insert("abc".into(), 1);
146 t.insert("def".into(), 2);
147
148 assert_eq!(t.remove("abc"), Some(1));
149
150 t.insert("ghi".into(), 3);
151 t.insert("jkl".into(), 4);
152
153 assert_eq!(t.get("abc"), None);
154 assert_eq!(t.get("def"), None);
155 assert_eq!(t.get("ghi"), Some(&3));
156 assert_eq!(t.get("jkl"), Some(&4));
157 }
158
159 #[test]
160 fn test_evicts_after_second_oldest_item_removed() {
161 let mut t = Test::new(3);
162 t.insert("abc".into(), 1);
163 t.insert("def".into(), 2);
164
165 assert_eq!(t.remove("def"), Some(2));
166 assert_eq!(t.get("abc"), Some(&1));
167
168 t.insert("ghi".into(), 3);
169 t.insert("jkl".into(), 4);
170
171 assert_eq!(t.get("abc"), None);
172 assert_eq!(t.get("def"), None);
173 assert_eq!(t.get("ghi"), Some(&3));
174 assert_eq!(t.get("jkl"), Some(&4));
175 }
176
177 #[test]
178 fn test_removes_all_items() {
179 let mut t = Test::new(3);
180 t.insert("abc".into(), 1);
181 t.insert("def".into(), 2);
182
183 assert_eq!(t.remove("def"), Some(2));
184 assert_eq!(t.remove("abc"), Some(1));
185
186 t.insert("ghi".into(), 3);
187 t.insert("jkl".into(), 4);
188 t.insert("mno".into(), 5);
189
190 assert_eq!(t.get("abc"), None);
191 assert_eq!(t.get("def"), None);
192 assert_eq!(t.get("ghi"), None);
193 assert_eq!(t.get("jkl"), Some(&4));
194 assert_eq!(t.get("mno"), Some(&5));
195 }
196
197 #[test]
198 fn test_inserts_many_items() {
199 let mut t = Test::new(3);
200
201 for _ in 0..10000 {
202 t.insert("abc".into(), 1);
203 t.insert("def".into(), 2);
204 t.insert("ghi".into(), 3);
205 }
206 }
207
208 #[test]
209 fn test_get_or_insert_default_and_edit_evicts_old_items_to_meet_capacity() {
210 let mut t = Test::new(3);
211
212 t.get_or_insert_default_and_edit("abc".into(), |v| *v += 1);
213 t.get_or_insert_default_and_edit("def".into(), |v| *v += 2);
214
215 // evicts "abc"
216 t.get_or_insert_default_and_edit("ghi".into(), |v| *v += 3);
217 assert_eq!(t.get("abc"), None);
218
219 // evicts "def"
220 t.get_or_insert_default_and_edit("jkl".into(), |v| *v += 4);
221 assert_eq!(t.get("def"), None);
222
223 // evicts "ghi"
224 t.get_or_insert_default_and_edit("abc".into(), |v| *v += 5);
225 assert_eq!(t.get("ghi"), None);
226
227 // evicts "jkl"
228 t.get_or_insert_default_and_edit("def".into(), |v| *v += 6);
229
230 assert_eq!(t.get("abc"), Some(&5));
231 assert_eq!(t.get("def"), Some(&6));
232 assert_eq!(t.get("ghi"), None);
233 assert_eq!(t.get("jkl"), None);
234 }
235
236 #[test]
237 fn test_get_or_insert_default_and_edit_edits_existing_item() {
238 let mut t = Test::new(3);
239
240 t.get_or_insert_default_and_edit("abc".into(), |v| *v += 1);
241 t.get_or_insert_default_and_edit("abc".into(), |v| *v += 2);
242 t.get_or_insert_default_and_edit("abc".into(), |v| *v += 3);
243
244 assert_eq!(t.get("abc"), Some(&6));
245 }
246}
247