arc_cache/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use std::borrow::Borrow;
4use std::collections::hash_map::RandomState;
5use std::hash::{BuildHasher, Hash};
6use std::iter::Chain;
7use xlru_cache::LruCache;
8
9pub struct ArcCache<K, V, S = RandomState>
10where
11    K: Eq + Hash,
12    S: BuildHasher,
13{
14    recent_set: LruCache<K, V, S>,
15    recent_evicted: LruCache<K, (), S>,
16    frequent_set: LruCache<K, V, S>,
17    frequent_evicted: LruCache<K, (), S>,
18    capacity: usize,
19    p: usize,
20    inserted: u64,
21    evicted: u64,
22    removed: u64,
23}
24
25/// An iterator over all items in the cache. Iterates over frequently-used items
26/// followed by recently-used items. Within each of these groups, the iterator
27/// yields less recently used items first.
28pub struct ArcCacheIterator<'a, K, V>(
29    Chain<xlru_cache::Iter<'a, K, V>, xlru_cache::Iter<'a, K, V>>,
30)
31where
32    K: Eq + Hash;
33
34impl<K, V> ArcCache<K, V>
35where
36    K: Eq + Hash,
37{
38    pub fn new(capacity: usize) -> Result<Self, &'static str> {
39        if capacity == 0 {
40            return Err("Cache length cannot be zero");
41        }
42        let cache = ArcCache {
43            recent_set: LruCache::new(capacity),
44            recent_evicted: LruCache::new(capacity),
45            frequent_set: LruCache::new(capacity),
46            frequent_evicted: LruCache::new(capacity),
47            capacity,
48            p: 0,
49            inserted: 0,
50            evicted: 0,
51            removed: 0,
52        };
53        Ok(cache)
54    }
55}
56
57impl<K, V, S> ArcCache<K, V, S>
58where
59    K: Eq + Hash,
60    S: BuildHasher,
61{
62    /// Creates an empty cache that can hold at most `capacity` items with the given hash builder.
63    pub fn with_hasher(capacity: usize, hash_builder: S) -> Result<Self, &'static str>
64    where
65        S: Clone,
66    {
67        if capacity == 0 {
68            return Err("Cache length cannot be zero");
69        }
70        let cache = ArcCache {
71            recent_set: LruCache::with_hasher(capacity, hash_builder.clone()),
72            recent_evicted: LruCache::with_hasher(capacity, hash_builder.clone()),
73            frequent_set: LruCache::with_hasher(capacity, hash_builder.clone()),
74            frequent_evicted: LruCache::with_hasher(capacity, hash_builder),
75            capacity,
76            p: 0,
77            inserted: 0,
78            evicted: 0,
79            removed: 0,
80        };
81        Ok(cache)
82    }
83
84    pub fn contains_key<Q: ?Sized>(&mut self, key: &Q) -> bool
85    where
86        K: Borrow<Q>,
87        Q: Hash + Eq,
88    {
89        self.frequent_set.contains_key(key) || self.recent_set.contains_key(key)
90    }
91
92    pub fn insert(&mut self, key: K, value: V) -> bool {
93        if self.frequent_set.contains_key(&key) {
94            self.frequent_set.insert(key, value);
95            return true;
96        }
97        if self.recent_set.contains_key(&key) {
98            self.recent_set.remove(&key);
99            self.frequent_set.insert(key, value);
100            return true;
101        }
102        if self.frequent_evicted.contains_key(&key) {
103            let recent_evicted_len = self.recent_evicted.len();
104            let frequent_evicted_len = self.frequent_evicted.len();
105            let delta = if recent_evicted_len > frequent_evicted_len {
106                recent_evicted_len / frequent_evicted_len
107            } else {
108                1
109            };
110            if delta < self.p {
111                self.p -= delta;
112            } else {
113                self.p = 0
114            }
115            if self.recent_set.len() + self.frequent_set.len() >= self.capacity {
116                self.replace(true);
117            }
118            self.frequent_evicted.remove(&key);
119            self.frequent_set.insert(key, value);
120            return true;
121        }
122        if self.recent_evicted.contains_key(&key) {
123            let recent_evicted_len = self.recent_evicted.len();
124            let frequent_evicted_len = self.frequent_evicted.len();
125            let delta = if frequent_evicted_len > recent_evicted_len {
126                frequent_evicted_len / recent_evicted_len
127            } else {
128                1
129            };
130            if delta <= self.capacity - self.p {
131                self.p += delta;
132            } else {
133                self.p = self.capacity;
134            }
135            if self.recent_set.len() + self.frequent_set.len() >= self.capacity {
136                self.replace(false);
137            }
138            self.recent_evicted.remove(&key);
139            self.frequent_set.insert(key, value);
140            return true;
141        }
142        if self.recent_set.len() + self.frequent_set.len() >= self.capacity {
143            self.replace(false);
144        }
145        if self.recent_evicted.len() > self.capacity - self.p {
146            self.recent_evicted.remove_lru();
147            self.evicted += 1;
148        }
149        if self.frequent_evicted.len() > self.p {
150            self.frequent_evicted.remove_lru();
151            self.evicted += 1;
152        }
153        self.recent_set.insert(key, value);
154        self.inserted += 1;
155        false
156    }
157
158    pub fn peek_mut(&mut self, key: &K) -> Option<&mut V>
159    where
160        K: Clone + Hash + Eq,
161    {
162        if let Some(entry) = self.frequent_set.peek_mut(key) {
163            Some(entry)
164        } else {
165            self.recent_set.peek_mut(key)
166        }
167    }
168
169    pub fn get_mut(&mut self, key: &K) -> Option<&mut V>
170    where
171        K: Clone + Hash + Eq,
172    {
173        if let Some(value) = self.recent_set.remove(key) {
174            self.frequent_set.insert((*key).clone(), value);
175        }
176        self.frequent_set.get_mut(key)
177    }
178
179    pub fn remove(&mut self, key: &K) -> Option<V> {
180        let removed_frequent = self.frequent_set.remove(key);
181        let removed_recent = self.recent_set.remove(key);
182
183        self.frequent_evicted.remove(key);
184        self.recent_evicted.remove(key);
185
186        match removed_frequent.or(removed_recent) {
187            Some(value) => {
188                self.removed += 1;
189                Some(value)
190            }
191            None => None,
192        }
193    }
194
195    fn replace(&mut self, frequent_evicted_contains_key: bool) {
196        let recent_set_len = self.recent_set.len();
197        if recent_set_len > 0
198            && (recent_set_len > self.p
199                || (recent_set_len == self.p && frequent_evicted_contains_key))
200        {
201            if let Some((old_key, _)) = self.recent_set.remove_lru() {
202                self.recent_evicted.insert(old_key, ());
203            }
204        } else if let Some((old_key, _)) = self.frequent_set.remove_lru() {
205            self.frequent_evicted.insert(old_key, ());
206        }
207    }
208
209    pub fn clear(&mut self) {
210        self.recent_set.clear();
211        self.recent_evicted.clear();
212        self.frequent_set.clear();
213        self.frequent_evicted.clear();
214    }
215
216    pub fn len(&self) -> usize {
217        self.recent_set.len() + self.frequent_set.len()
218    }
219
220    pub fn is_empty(&self) -> bool {
221        self.len() == 0
222    }
223
224    pub fn frequent_len(&self) -> usize {
225        self.frequent_set.len()
226    }
227
228    pub fn recent_len(&self) -> usize {
229        self.recent_set.len()
230    }
231
232    pub fn inserted(&self) -> u64 {
233        self.inserted
234    }
235
236    pub fn evicted(&self) -> u64 {
237        self.evicted
238    }
239
240    pub fn removed(&self) -> u64 {
241        self.removed
242    }
243}
244
245impl<'a, K, V, S> IntoIterator for &'a ArcCache<K, V, S>
246where
247    K: Eq + Hash,
248    S: BuildHasher,
249{
250    type Item = (&'a K, &'a V);
251    type IntoIter = ArcCacheIterator<'a, K, V>;
252
253    fn into_iter(self) -> Self::IntoIter {
254        ArcCacheIterator(self.frequent_set.iter().chain(self.recent_set.iter()))
255    }
256}
257
258impl<'a, K, V> Iterator for ArcCacheIterator<'a, K, V>
259where
260    K: Eq + Hash,
261{
262    type Item = (&'a K, &'a V);
263
264    fn next(&mut self) -> Option<Self::Item> {
265        self.0.next()
266    }
267}
268
269#[test]
270fn test_arc() {
271    let mut arc: ArcCache<&str, &str> = ArcCache::new(2).unwrap();
272    arc.insert("testkey", "testvalue");
273    assert!(arc.contains_key(&"testkey"));
274    arc.insert("testkey2", "testvalue2");
275    assert!(arc.contains_key(&"testkey2"));
276    arc.insert("testkey3", "testvalue3");
277    assert!(arc.contains_key(&"testkey3"));
278    assert!(arc.contains_key(&"testkey2"));
279    assert!(!arc.contains_key(&"testkey"));
280    arc.insert("testkey", "testvalue");
281    assert!(arc.get_mut(&"testkey").is_some());
282    assert!(arc.get_mut(&"testkey-nx").is_none());
283
284    let mut it = arc.into_iter();
285    assert_eq!(it.next(), Some((&"testkey", &"testvalue")));
286    assert_eq!(it.next(), Some((&"testkey2", &"testvalue2")));
287    assert_eq!(it.next(), None);
288}
289
290#[test]
291fn test_custom_hasher() {
292    use nohash_hasher::NoHashHasher;
293    use std::hash::BuildHasherDefault;
294
295    let mut arc: ArcCache<u32, &str, BuildHasherDefault<NoHashHasher<u8>>> =
296        ArcCache::with_hasher(2, BuildHasherDefault::default()).unwrap();
297    arc.insert(1, "testvalue");
298    assert!(arc.contains_key(&1));
299    arc.insert(2, "testvalue2");
300    assert!(arc.contains_key(&2));
301    arc.insert(3, "testvalue3");
302    assert!(arc.contains_key(&3));
303    assert!(arc.contains_key(&2));
304    assert!(!arc.contains_key(&1));
305    arc.insert(1, "testvalue");
306    assert!(arc.get_mut(&1).is_some());
307    assert!(arc.get_mut(&666).is_none());
308
309    let mut it = arc.into_iter();
310    assert_eq!(it.next(), Some((&1, &"testvalue")));
311    assert_eq!(it.next(), Some((&2, &"testvalue2")));
312    assert_eq!(it.next(), None);
313}