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