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

rustls/
limited_cache.rs
