rustls/
limited_cache.rs

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

rustls/
limited_cache.rs
