rustls/
limited_cache.rs

Help
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}
rustls/limited_cache.rs

rustls/
limited_cache.rs
