rustls/
limited_cache.rs

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        if let Some(value) = self.map.remove(k) {
106            // O(N) search, followed by O(N) removal
107            if let Some(index) = self
108                .oldest
109                .iter()
110                .position(|item| item.borrow() == k)
111            {
112                self.oldest.remove(index);
113            }
114            Some(value)
115        } else {
116            None
117        }
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use std::prelude::v1::*;
124
125    type Test = super::LimitedCache<String, usize>;
126
127    #[test]
128    fn test_updates_existing_item() {
129        let mut t = Test::new(3);
130        t.insert("abc".into(), 1);
131        t.insert("abc".into(), 2);
132        assert_eq!(t.get("abc"), Some(&2));
133    }
134
135    #[test]
136    fn test_evicts_oldest_item() {
137        let mut t = Test::new(3);
138        t.insert("abc".into(), 1);
139        t.insert("def".into(), 2);
140        t.insert("ghi".into(), 3);
141
142        assert_eq!(t.get("abc"), None);
143        assert_eq!(t.get("def"), Some(&2));
144        assert_eq!(t.get("ghi"), Some(&3));
145    }
146
147    #[test]
148    fn test_evicts_second_oldest_item_if_first_removed() {
149        let mut t = Test::new(3);
150        t.insert("abc".into(), 1);
151        t.insert("def".into(), 2);
152
153        assert_eq!(t.remove("abc"), Some(1));
154
155        t.insert("ghi".into(), 3);
156        t.insert("jkl".into(), 4);
157
158        assert_eq!(t.get("abc"), None);
159        assert_eq!(t.get("def"), None);
160        assert_eq!(t.get("ghi"), Some(&3));
161        assert_eq!(t.get("jkl"), Some(&4));
162    }
163
164    #[test]
165    fn test_evicts_after_second_oldest_item_removed() {
166        let mut t = Test::new(3);
167        t.insert("abc".into(), 1);
168        t.insert("def".into(), 2);
169
170        assert_eq!(t.remove("def"), Some(2));
171        assert_eq!(t.get("abc"), Some(&1));
172
173        t.insert("ghi".into(), 3);
174        t.insert("jkl".into(), 4);
175
176        assert_eq!(t.get("abc"), None);
177        assert_eq!(t.get("def"), None);
178        assert_eq!(t.get("ghi"), Some(&3));
179        assert_eq!(t.get("jkl"), Some(&4));
180    }
181
182    #[test]
183    fn test_removes_all_items() {
184        let mut t = Test::new(3);
185        t.insert("abc".into(), 1);
186        t.insert("def".into(), 2);
187
188        assert_eq!(t.remove("def"), Some(2));
189        assert_eq!(t.remove("abc"), Some(1));
190
191        t.insert("ghi".into(), 3);
192        t.insert("jkl".into(), 4);
193        t.insert("mno".into(), 5);
194
195        assert_eq!(t.get("abc"), None);
196        assert_eq!(t.get("def"), None);
197        assert_eq!(t.get("ghi"), None);
198        assert_eq!(t.get("jkl"), Some(&4));
199        assert_eq!(t.get("mno"), Some(&5));
200    }
201
202    #[test]
203    fn test_inserts_many_items() {
204        let mut t = Test::new(3);
205
206        for _ in 0..10000 {
207            t.insert("abc".into(), 1);
208            t.insert("def".into(), 2);
209            t.insert("ghi".into(), 3);
210        }
211    }
212
213    #[test]
214    fn test_get_or_insert_default_and_edit_evicts_old_items_to_meet_capacity() {
215        let mut t = Test::new(3);
216
217        t.get_or_insert_default_and_edit("abc".into(), |v| *v += 1);
218        t.get_or_insert_default_and_edit("def".into(), |v| *v += 2);
219
220        // evicts "abc"
221        t.get_or_insert_default_and_edit("ghi".into(), |v| *v += 3);
222        assert_eq!(t.get("abc"), None);
223
224        // evicts "def"
225        t.get_or_insert_default_and_edit("jkl".into(), |v| *v += 4);
226        assert_eq!(t.get("def"), None);
227
228        // evicts "ghi"
229        t.get_or_insert_default_and_edit("abc".into(), |v| *v += 5);
230        assert_eq!(t.get("ghi"), None);
231
232        // evicts "jkl"
233        t.get_or_insert_default_and_edit("def".into(), |v| *v += 6);
234
235        assert_eq!(t.get("abc"), Some(&5));
236        assert_eq!(t.get("def"), Some(&6));
237        assert_eq!(t.get("ghi"), None);
238        assert_eq!(t.get("jkl"), None);
239    }
240
241    #[test]
242    fn test_get_or_insert_default_and_edit_edits_existing_item() {
243        let mut t = Test::new(3);
244
245        t.get_or_insert_default_and_edit("abc".into(), |v| *v += 1);
246        t.get_or_insert_default_and_edit("abc".into(), |v| *v += 2);
247        t.get_or_insert_default_and_edit("abc".into(), |v| *v += 3);
248
249        assert_eq!(t.get("abc"), Some(&6));
250    }
251}
rustls/limited_cache.rs

rustls/
limited_cache.rs
