Skip to main content

quick_cache/
unsync.rs

1use crate::{
2    linked_slab::Token,
3    options::*,
4    shard::{self, CacheShard, InsertStrategy},
5    DefaultHashBuilder, Equivalent, Lifecycle, MemoryUsed, UnitWeighter, Weighter,
6};
7use std::hash::{BuildHasher, Hash};
8
9/// A non-concurrent cache.
10#[derive(Clone)]
11pub struct Cache<
12    Key,
13    Val,
14    We = UnitWeighter,
15    B = DefaultHashBuilder,
16    L = DefaultLifecycle<Key, Val>,
17> {
18    shard: CacheShard<Key, Val, We, B, L, SharedPlaceholder>,
19}
20
21impl<Key: Eq + Hash, Val> Cache<Key, Val> {
22    /// Creates a new cache with holds up to `items_capacity` items (approximately).
23    pub fn new(items_capacity: usize) -> Self {
24        Self::with(
25            items_capacity,
26            items_capacity as u64,
27            Default::default(),
28            Default::default(),
29            Default::default(),
30        )
31    }
32}
33
34impl<Key: Eq + Hash, Val, We: Weighter<Key, Val>> Cache<Key, Val, We> {
35    pub fn with_weighter(
36        estimated_items_capacity: usize,
37        weight_capacity: u64,
38        weighter: We,
39    ) -> Self {
40        Self::with(
41            estimated_items_capacity,
42            weight_capacity,
43            weighter,
44            Default::default(),
45            Default::default(),
46        )
47    }
48}
49
50impl<Key: Eq + Hash, Val, We: Weighter<Key, Val>, B: BuildHasher, L: Lifecycle<Key, Val>>
51    Cache<Key, Val, We, B, L>
52{
53    /// Creates a new cache that can hold up to `weight_capacity` in weight.
54    /// `estimated_items_capacity` is the estimated number of items the cache is expected to hold,
55    /// roughly equivalent to `weight_capacity / average item weight`.
56    pub fn with(
57        estimated_items_capacity: usize,
58        weight_capacity: u64,
59        weighter: We,
60        hash_builder: B,
61        lifecycle: L,
62    ) -> Self {
63        Self::with_options(
64            OptionsBuilder::new()
65                .estimated_items_capacity(estimated_items_capacity)
66                .weight_capacity(weight_capacity)
67                .build()
68                .unwrap(),
69            weighter,
70            hash_builder,
71            lifecycle,
72        )
73    }
74
75    /// Constructs a cache based on [OptionsBuilder].
76    ///
77    /// # Example
78    ///
79    /// ```rust
80    /// use quick_cache::{unsync::{Cache, DefaultLifecycle}, OptionsBuilder, UnitWeighter, DefaultHashBuilder};
81    ///
82    /// Cache::<(String, u64), String>::with_options(
83    ///   OptionsBuilder::new()
84    ///     .estimated_items_capacity(10000)
85    ///     .weight_capacity(10000)
86    ///     .build()
87    ///     .unwrap(),
88    ///     UnitWeighter,
89    ///     DefaultHashBuilder::default(),
90    ///     DefaultLifecycle::default(),
91    /// );
92    /// ```
93    pub fn with_options(options: Options, weighter: We, hash_builder: B, lifecycle: L) -> Self {
94        let shard = CacheShard::new(
95            options.hot_allocation,
96            options.ghost_allocation,
97            options.estimated_items_capacity,
98            options.weight_capacity,
99            weighter,
100            hash_builder,
101            lifecycle,
102        );
103        Self { shard }
104    }
105
106    /// Returns whether the cache is empty.
107    pub fn is_empty(&self) -> bool {
108        self.shard.len() == 0
109    }
110
111    /// Returns the number of cached items
112    pub fn len(&self) -> usize {
113        self.shard.len()
114    }
115
116    /// Returns the total weight of cached items
117    pub fn weight(&self) -> u64 {
118        self.shard.weight()
119    }
120
121    /// Returns the maximum weight of cached items
122    pub fn capacity(&self) -> u64 {
123        self.shard.capacity()
124    }
125
126    /// Returns the number of misses
127    #[cfg(feature = "stats")]
128    pub fn misses(&self) -> u64 {
129        self.shard.misses()
130    }
131
132    /// Returns the number of hits
133    #[cfg(feature = "stats")]
134    pub fn hits(&self) -> u64 {
135        self.shard.hits()
136    }
137
138    /// Reserve additional space for `additional` entries.
139    /// Note that this is counted in entries, and is not weighted.
140    pub fn reserve(&mut self, additional: usize) {
141        self.shard.reserve(additional);
142    }
143
144    /// Checks if a key exists in the cache.
145    pub fn contains_key<Q>(&self, key: &Q) -> bool
146    where
147        Q: Hash + Equivalent<Key> + ?Sized,
148    {
149        self.shard.contains(self.shard.hash(key), key)
150    }
151
152    /// Fetches an item from the cache.
153    pub fn get<Q>(&self, key: &Q) -> Option<&Val>
154    where
155        Q: Hash + Equivalent<Key> + ?Sized,
156    {
157        self.shard.get(self.shard.hash(key), key)
158    }
159
160    /// Fetches an item from the cache.
161    ///
162    /// Note: Leaking the returned RefMut might cause cache weight tracking to be inaccurate.
163    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<RefMut<'_, Key, Val, We, B, L>>
164    where
165        Q: Hash + Equivalent<Key> + ?Sized,
166    {
167        self.shard.get_mut(self.shard.hash(key), key).map(RefMut)
168    }
169
170    /// Peeks an item from the cache. Contrary to gets, peeks don't alter the key "hotness".
171    pub fn peek<Q>(&self, key: &Q) -> Option<&Val>
172    where
173        Q: Hash + Equivalent<Key> + ?Sized,
174    {
175        self.shard.peek(self.shard.hash(key), key)
176    }
177
178    /// Returns per-item statistics for `key`, or `None` if the key is not present.
179    /// Like peeks, this does not alter the key "hotness" or its access count.
180    #[cfg(feature = "stats")]
181    pub fn item_stats<Q>(&self, key: &Q) -> Option<crate::ItemStats>
182    where
183        Q: Hash + Equivalent<Key> + ?Sized,
184    {
185        self.shard.item_stats(self.shard.hash(key), key)
186    }
187
188    /// Peeks an item from the cache. Contrary to gets, peeks don't alter the key "hotness".
189    ///
190    /// Note: Leaking the returned RefMut might cause cache weight tracking to be inaccurate.
191    pub fn peek_mut<Q>(&mut self, key: &Q) -> Option<RefMut<'_, Key, Val, We, B, L>>
192    where
193        Q: Hash + Equivalent<Key> + ?Sized,
194    {
195        self.shard.peek_mut(self.shard.hash(key), key).map(RefMut)
196    }
197
198    /// Remove an item from the cache whose key is `key`.
199    /// Returns the removed entry, if any.
200    pub fn remove<Q>(&mut self, key: &Q) -> Option<(Key, Val)>
201    where
202        Q: Hash + Equivalent<Key> + ?Sized,
203    {
204        self.shard.remove(self.shard.hash(key), key)
205    }
206
207    /// Remove an item from the cache whose key is `key` if `f(&value)` returns `true` for that entry.
208    /// Compared to peek and remove, this method is more efficient as it requires only 1 lookup.
209    ///
210    /// Returns the removed entry, if any.
211    pub fn remove_if<Q, F>(&mut self, key: &Q, f: F) -> Option<(Key, Val)>
212    where
213        Q: Hash + Equivalent<Key> + ?Sized,
214        F: FnOnce(&Val) -> bool,
215    {
216        self.shard.remove_if(self.shard.hash(key), key, f)
217    }
218
219    /// Replaces an item in the cache, but only if it already exists.
220    /// If `soft` is set, the replace operation won't affect the "hotness" of the key,
221    /// even if the value is replaced.
222    ///
223    /// Returns `Ok` if the entry was admitted and `Err(_)` if it wasn't.
224    pub fn replace(&mut self, key: Key, value: Val, soft: bool) -> Result<(), (Key, Val)> {
225        let mut lcs = Default::default();
226        self.replace_with_lifecycle(key, value, soft, &mut lcs)
227    }
228
229    /// Replaces an item in the cache, but only if it already exists, recording any evicted
230    /// items into the given lifecycle request state.
231    /// If `soft` is set, the replace operation won't affect the "hotness" of the key,
232    /// even if the value is replaced.
233    ///
234    /// Returns `Ok` if the entry was admitted and `Err(_)` if it wasn't.
235    ///
236    /// `lcs` is a [`Lifecycle::RequestState`]; construct one with `Default::default()`.
237    /// The same `&mut lcs` can be threaded through multiple operations to batch the
238    /// eviction work. Evicted items are released when `lcs` is dropped.
239    pub fn replace_with_lifecycle(
240        &mut self,
241        key: Key,
242        value: Val,
243        soft: bool,
244        lcs: &mut L::RequestState,
245    ) -> Result<(), (Key, Val)> {
246        self.shard.insert(
247            lcs,
248            self.shard.hash(&key),
249            key,
250            value,
251            InsertStrategy::Replace { soft },
252        )?;
253        Ok(())
254    }
255
256    /// Retains only the items specified by the predicate.
257    /// In other words, remove all items for which `f(&key, &value)` returns `false`. The
258    /// elements are visited in arbitrary order.
259    pub fn retain<F>(&mut self, f: F)
260    where
261        F: Fn(&Key, &Val) -> bool,
262    {
263        self.shard.retain(f);
264    }
265
266    /// Gets or inserts an item in the cache with key `key`.
267    /// Returns a reference to the inserted `value` if it was admitted to the cache.
268    ///
269    /// See also `get_ref_or_guard`.
270    pub fn get_or_insert_with<Q, E>(
271        &mut self,
272        key: &Q,
273        with: impl FnOnce() -> Result<Val, E>,
274    ) -> Result<Option<&Val>, E>
275    where
276        Q: Hash + Equivalent<Key> + ToOwned<Owned = Key> + ?Sized,
277    {
278        let idx = match self.shard.get_or_placeholder(self.shard.hash(key), key) {
279            Ok((idx, _)) => idx,
280            Err((plh, _)) => {
281                let v = with()?;
282                let mut lcs = Default::default();
283                let replaced = self.shard.replace_placeholder(&mut lcs, &plh, false, v);
284                debug_assert!(replaced.is_ok(), "unsync replace_placeholder can't fail");
285                plh.idx
286            }
287        };
288        Ok(self.shard.peek_token(idx))
289    }
290
291    /// Gets or inserts an item in the cache with key `key`.
292    /// Returns a mutable reference to the inserted `value` if it was admitted to the cache.
293    ///
294    /// See also `get_mut_or_guard`.
295    pub fn get_mut_or_insert_with<'a, Q, E>(
296        &'a mut self,
297        key: &Q,
298        with: impl FnOnce() -> Result<Val, E>,
299    ) -> Result<Option<RefMut<'a, Key, Val, We, B, L>>, E>
300    where
301        Q: Hash + Equivalent<Key> + ToOwned<Owned = Key> + ?Sized,
302    {
303        let idx = match self.shard.get_or_placeholder(self.shard.hash(key), key) {
304            Ok((idx, _)) => idx,
305            Err((plh, _)) => {
306                let v = with()?;
307                let mut lcs = Default::default();
308                let replaced = self.shard.replace_placeholder(&mut lcs, &plh, false, v);
309                debug_assert!(replaced.is_ok(), "unsync replace_placeholder can't fail");
310                plh.idx
311            }
312        };
313        Ok(self.shard.peek_token_mut(idx).map(RefMut))
314    }
315
316    /// Gets an item from the cache with key `key` .
317    /// If the corresponding value isn't present in the cache, this function returns a guard
318    /// that can be used to insert the value once it's computed.
319    pub fn get_ref_or_guard<Q>(&mut self, key: &Q) -> Result<&Val, Guard<'_, Key, Val, We, B, L>>
320    where
321        Q: Hash + Equivalent<Key> + ToOwned<Owned = Key> + ?Sized,
322    {
323        // TODO: this could be using a simpler entry API
324        match self.shard.get_or_placeholder(self.shard.hash(key), key) {
325            Ok((_, v)) => unsafe {
326                // Rustc gets insanely confused about returning from mut borrows
327                // Safety: v has the same lifetime as self
328                let v: *const Val = v;
329                Ok(&*v)
330            },
331            Err((placeholder, _)) => Err(Guard {
332                cache: self,
333                placeholder,
334                inserted: false,
335            }),
336        }
337    }
338
339    /// Gets an item from the cache with key `key` .
340    /// If the corresponding value isn't present in the cache, this function returns a guard
341    /// that can be used to insert the value once it's computed.
342    ///
343    /// Note: Leaking the returned RefMut might cause cache weight tracking to be inaccurate.
344    pub fn get_mut_or_guard<'a, Q>(
345        &'a mut self,
346        key: &Q,
347    ) -> Result<Option<RefMut<'a, Key, Val, We, B, L>>, Guard<'a, Key, Val, We, B, L>>
348    where
349        Q: Hash + Equivalent<Key> + ToOwned<Owned = Key> + ?Sized,
350    {
351        // TODO: this could be using a simpler entry API
352        match self.shard.get_or_placeholder(self.shard.hash(key), key) {
353            Ok((idx, _)) => Ok(self.shard.peek_token_mut(idx).map(RefMut)),
354            Err((placeholder, _)) => Err(Guard {
355                cache: self,
356                placeholder,
357                inserted: false,
358            }),
359        }
360    }
361
362    /// Inserts an item in the cache with key `key`.
363    pub fn insert(&mut self, key: Key, value: Val) {
364        let mut lcs = Default::default();
365        self.insert_with_lifecycle(key, value, &mut lcs);
366    }
367
368    /// Inserts an item in the cache with key `key`, recording any evicted items into the
369    /// given lifecycle request state.
370    ///
371    /// `lcs` is a [`Lifecycle::RequestState`]; construct one with `Default::default()`.
372    /// The same `&mut lcs` can be threaded through multiple operations to batch the
373    /// eviction work. Evicted items are released when `lcs` is dropped.
374    pub fn insert_with_lifecycle(&mut self, key: Key, value: Val, lcs: &mut L::RequestState) {
375        let result = self.shard.insert(
376            lcs,
377            self.shard.hash(&key),
378            key,
379            value,
380            InsertStrategy::Insert,
381        );
382        // result cannot err with the Insert strategy
383        debug_assert!(result.is_ok());
384    }
385
386    /// Clear all items from the cache
387    pub fn clear(&mut self) {
388        self.shard.clear();
389    }
390
391    /// Iterator for the items in the cache
392    pub fn iter(&self) -> impl Iterator<Item = (&'_ Key, &'_ Val)> + '_ {
393        // TODO: add a concrete type, impl trait in the public api is really bad.
394        self.shard.iter()
395    }
396
397    /// Drain all items from the cache
398    ///
399    /// The cache will be emptied even if the returned iterator isn't fully consumed.
400    pub fn drain(&mut self) -> impl Iterator<Item = (Key, Val)> + '_ {
401        // TODO: add a concrete type, impl trait in the public api is really bad.
402        self.shard.drain()
403    }
404
405    /// Sets the cache to a new weight capacity.
406    ///
407    /// If the new capacity is smaller than the current weight, items will be evicted
408    /// to bring the cache within the new limit.
409    pub fn set_capacity(&mut self, new_weight_capacity: u64) {
410        let mut lcs = Default::default();
411        self.shard.set_capacity(new_weight_capacity, &mut lcs);
412    }
413
414    #[cfg(any(fuzzing, test))]
415    pub fn validate(&self, accept_overweight: bool) {
416        self.shard.validate(accept_overweight);
417    }
418
419    /// Get total memory used by cache data structures
420    ///
421    /// It should be noted that if cache key or value is some type like `Vec<T>`,
422    /// the memory allocated in the heap will not be counted.
423    pub fn memory_used(&self) -> MemoryUsed {
424        self.shard.memory_used()
425    }
426}
427
428impl<Key, Val, We, B, L> std::fmt::Debug for Cache<Key, Val, We, B, L> {
429    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
430        f.debug_struct("Cache").finish_non_exhaustive()
431    }
432}
433
434/// Default `Lifecycle` for the unsync cache.
435pub struct DefaultLifecycle<Key, Val>(std::marker::PhantomData<(Key, Val)>);
436
437impl<Key, Val> std::fmt::Debug for DefaultLifecycle<Key, Val> {
438    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
439        f.debug_tuple("DefaultLifecycle").finish()
440    }
441}
442
443impl<Key, Val> Default for DefaultLifecycle<Key, Val> {
444    #[inline]
445    fn default() -> Self {
446        Self(Default::default())
447    }
448}
449
450impl<Key, Val> Clone for DefaultLifecycle<Key, Val> {
451    #[inline]
452    fn clone(&self) -> Self {
453        Self(Default::default())
454    }
455}
456
457impl<Key, Val> Lifecycle<Key, Val> for DefaultLifecycle<Key, Val> {
458    type RequestState = ();
459}
460
461#[derive(Debug, Clone)]
462pub(crate) struct SharedPlaceholder {
463    hash: u64,
464    idx: Token,
465}
466
467pub struct Guard<'a, Key, Val, We, B, L> {
468    cache: &'a mut Cache<Key, Val, We, B, L>,
469    placeholder: SharedPlaceholder,
470    inserted: bool,
471}
472
473impl<Key: Eq + Hash, Val, We: Weighter<Key, Val>, B: BuildHasher, L: Lifecycle<Key, Val>>
474    Guard<'_, Key, Val, We, B, L>
475{
476    /// Inserts the value into the placeholder
477    pub fn insert(self, value: Val) {
478        let mut lcs = Default::default();
479        self.insert_with_lifecycle(value, &mut lcs);
480    }
481
482    /// Inserts the value into the placeholder, recording any evicted items into the given
483    /// lifecycle request state.
484    ///
485    /// `lcs` is a [`Lifecycle::RequestState`]; construct one with `Default::default()`.
486    /// Evicted items are released when `lcs` is dropped.
487    pub fn insert_with_lifecycle(mut self, value: Val, lcs: &mut L::RequestState) {
488        let replaced = self
489            .cache
490            .shard
491            .replace_placeholder(lcs, &self.placeholder, false, value);
492        debug_assert!(replaced.is_ok(), "unsync replace_placeholder can't fail");
493        self.inserted = true;
494    }
495}
496
497impl<Key, Val, We, B, L> Drop for Guard<'_, Key, Val, We, B, L> {
498    #[inline]
499    fn drop(&mut self) {
500        #[cold]
501        fn drop_slow<Key, Val, We, B, L>(this: &mut Guard<'_, Key, Val, We, B, L>) {
502            this.cache.shard.remove_placeholder(&this.placeholder);
503        }
504        if !self.inserted {
505            drop_slow(self);
506        }
507    }
508}
509
510pub struct RefMut<'cache, Key, Val, We: Weighter<Key, Val>, B, L>(
511    crate::shard::RefMut<'cache, Key, Val, We, B, L, SharedPlaceholder>,
512);
513
514impl<Key, Val, We: Weighter<Key, Val>, B, L> std::ops::Deref for RefMut<'_, Key, Val, We, B, L> {
515    type Target = Val;
516
517    #[inline]
518    fn deref(&self) -> &Self::Target {
519        self.0.pair().1
520    }
521}
522
523impl<Key, Val, We: Weighter<Key, Val>, B, L> std::ops::DerefMut for RefMut<'_, Key, Val, We, B, L> {
524    #[inline]
525    fn deref_mut(&mut self) -> &mut Self::Target {
526        self.0.value_mut()
527    }
528}
529
530impl shard::SharedPlaceholder for SharedPlaceholder {
531    #[inline]
532    fn new(hash: u64, idx: Token) -> Self {
533        Self { hash, idx }
534    }
535
536    #[inline]
537    fn same_as(&self, _other: &Self) -> bool {
538        true
539    }
540
541    #[inline]
542    fn hash(&self) -> u64 {
543        self.hash
544    }
545
546    #[inline]
547    fn idx(&self) -> Token {
548        self.idx
549    }
550}
551
552#[cfg(test)]
553mod tests {
554    use super::*;
555
556    struct Weighter;
557
558    impl crate::Weighter<u32, u32> for Weighter {
559        fn weight(&self, _key: &u32, val: &u32) -> u64 {
560            *val as u64
561        }
562    }
563
564    #[test]
565    fn test_zero_weights() {
566        let mut cache = Cache::with_weighter(100, 100, Weighter);
567        cache.insert(0, 0);
568        assert_eq!(cache.weight(), 0);
569        for i in 1..100 {
570            cache.insert(i, i);
571            cache.insert(i, i);
572        }
573        assert_eq!(cache.get(&0).copied(), Some(0));
574        assert!(cache.contains_key(&0));
575        let a = cache.weight();
576        *cache.get_mut(&0).unwrap() += 1;
577        assert_eq!(cache.weight(), a + 1);
578        for i in 1..100 {
579            cache.insert(i, i);
580            cache.insert(i, i);
581        }
582        assert_eq!(cache.get(&0), None);
583        assert!(!cache.contains_key(&0));
584
585        cache.insert(0, 1);
586        let a = cache.weight();
587        *cache.get_mut(&0).unwrap() -= 1;
588        assert_eq!(cache.weight(), a - 1);
589        for i in 1..100 {
590            cache.insert(i, i);
591            cache.insert(i, i);
592        }
593        assert_eq!(cache.get(&0).copied(), Some(0));
594        assert!(cache.contains_key(&0));
595    }
596
597    #[test]
598    fn test_set_capacity() {
599        let mut cache = Cache::new(100);
600        for i in 0..80 {
601            cache.insert(i, i);
602        }
603        let initial_len = cache.len();
604        assert!(initial_len <= 80);
605
606        // Set to smaller capacity
607        cache.set_capacity(50);
608        assert!(cache.len() <= 50);
609        assert!(cache.weight() <= 50);
610        cache.validate(false);
611
612        // Set to larger capacity
613        cache.set_capacity(200);
614        assert_eq!(cache.capacity(), 200);
615        cache.validate(false);
616
617        // Insert more items
618        for i in 100..180 {
619            cache.insert(i, i);
620        }
621        assert!(cache.len() <= 180);
622        assert!(cache.weight() <= 200);
623        cache.validate(false);
624    }
625
626    #[test]
627    fn test_set_capacity_with_ghosts() {
628        // Create a cache that will generate ghost entries
629        let mut cache = Cache::new(50);
630
631        // Insert items to fill the cache
632        for i in 0..100 {
633            cache.insert(i, i);
634        }
635        cache.validate(false);
636
637        // Set to smaller capacity - should trim both resident and ghost entries
638        cache.set_capacity(25);
639        assert!(cache.weight() <= 25);
640        cache.validate(false);
641
642        // Set back to larger capacity
643        cache.set_capacity(100);
644        assert_eq!(cache.capacity(), 100);
645        cache.validate(false);
646
647        // Insert more items
648        for i in 100..150 {
649            cache.insert(i, i);
650        }
651        cache.validate(false);
652    }
653
654    #[test]
655    fn test_remove_if() {
656        let mut cache = Cache::new(100);
657
658        // Insert test data
659        cache.insert(1, 10);
660        cache.insert(2, 20);
661        cache.insert(3, 30);
662
663        // Test removing with predicate that returns true
664        let removed = cache.remove_if(&2, |v| *v == 20);
665        assert_eq!(removed, Some((2, 20)));
666        assert_eq!(cache.get(&2), None);
667
668        // Test removing with predicate that returns false
669        let not_removed = cache.remove_if(&3, |v| *v == 999);
670        assert_eq!(not_removed, None);
671        assert_eq!(cache.get(&3), Some(&30));
672
673        // Test removing non-existent key
674        let not_found = cache.remove_if(&999, |_| true);
675        assert_eq!(not_found, None);
676
677        cache.validate(false);
678    }
679}