Skip to main content

micro_moka/unsync/
cache.rs

1use super::{CacheBuilder, IndexDeque, Iter, Slab, SlabEntry, SENTINEL};
2use crate::{
3    common::{self, frequency_sketch::FrequencySketch},
4    Policy,
5};
6
7use hashbrown::HashTable;
8use std::{
9    borrow::Borrow,
10    collections::hash_map::RandomState,
11    fmt,
12    hash::{BuildHasher, Hash},
13};
14
15const EVICTION_BATCH_SIZE: usize = 100;
16
17/// An in-memory cache that is _not_ thread-safe.
18///
19/// `Cache` utilizes a hash table [`hashbrown::HashTable`][hb-hashtable] for the
20/// central key-value storage. `Cache` performs a best-effort bounding of the
21/// map using an entry replacement algorithm to determine which entries to evict
22/// when the capacity is exceeded.
23///
24/// [hb-hashtable]: https://docs.rs/hashbrown/latest/hashbrown/struct.HashTable.html
25///
26/// # Characteristic difference between `unsync` and `sync`/`future` caches
27///
28/// If you use a cache from a single thread application, `unsync::Cache` may
29/// outperform other caches for updates and retrievals because other caches have some
30/// overhead on syncing internal data structures between threads.
31///
32/// # Examples
33///
34/// Cache entries are manually added using the insert method, and are stored in the
35/// cache until either evicted or manually invalidated.
36///
37/// Here's an example of reading and updating a cache by using the main thread:
38///
39///```rust
40/// use micro_moka::unsync::Cache;
41///
42/// const NUM_KEYS: usize = 64;
43///
44/// fn value(n: usize) -> String {
45///     format!("value {}", n)
46/// }
47///
48/// // Create a cache that can store up to 10,000 entries.
49/// let mut cache = Cache::new(10_000);
50///
51/// // Insert 64 entries.
52/// for key in 0..NUM_KEYS {
53///     cache.insert(key, value(key));
54/// }
55///
56/// // Invalidate every 4 element of the inserted entries.
57/// for key in (0..NUM_KEYS).step_by(4) {
58///     cache.invalidate(&key);
59/// }
60///
61/// // Verify the result.
62/// for key in 0..NUM_KEYS {
63///     if key % 4 == 0 {
64///         assert_eq!(cache.get(&key), None);
65///     } else {
66///         assert_eq!(cache.get(&key), Some(&value(key)));
67///     }
68/// }
69/// ```
70///
71/// # Hashing Algorithm
72///
73/// By default, `Cache` uses a hashing algorithm selected to provide resistance
74/// against HashDoS attacks. It will the same one used by
75/// `std::collections::HashMap`, which is currently SipHash 1-3.
76///
77/// While SipHash's performance is very competitive for medium sized keys, other
78/// hashing algorithms will outperform it for small keys such as integers as well as
79/// large keys such as long strings. However those algorithms will typically not
80/// protect against attacks such as HashDoS.
81///
82/// The hashing algorithm can be replaced on a per-`Cache` basis using the
83/// [`build_with_hasher`][build-with-hasher-method] method of the
84/// `CacheBuilder`. Many alternative algorithms are available on crates.io, such
85/// as the [aHash][ahash-crate] crate.
86///
87/// [build-with-hasher-method]: ./struct.CacheBuilder.html#method.build_with_hasher
88/// [ahash-crate]: https://crates.io/crates/ahash
89///
90pub struct Cache<K, V, S = RandomState> {
91    max_capacity: Option<u64>,
92    entry_count: u64,
93    table: HashTable<u32>,
94    build_hasher: S,
95    slab: Slab<K, V>,
96    deque: IndexDeque,
97    frequency_sketch: FrequencySketch,
98    frequency_sketch_enabled: bool,
99}
100
101impl<K, V, S> fmt::Debug for Cache<K, V, S>
102where
103    K: fmt::Debug + Eq + Hash,
104    V: fmt::Debug,
105    S: BuildHasher + Clone,
106{
107    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
108        let mut d_map = f.debug_map();
109
110        for (k, v) in self.iter() {
111            d_map.entry(&k, &v);
112        }
113
114        d_map.finish()
115    }
116}
117
118impl<K, V> Cache<K, V, RandomState>
119where
120    K: Hash + Eq,
121{
122    /// Constructs a new `Cache<K, V>` that will store up to the `max_capacity` entries.
123    ///
124    /// To adjust various configuration knobs such as `initial_capacity`, use the
125    /// [`CacheBuilder`][builder-struct].
126    ///
127    /// [builder-struct]: ./struct.CacheBuilder.html
128    pub fn new(max_capacity: u64) -> Self {
129        let build_hasher = RandomState::default();
130        Self::with_everything(Some(max_capacity), None, build_hasher)
131    }
132
133    /// Returns a [`CacheBuilder`][builder-struct], which can builds a `Cache` with
134    /// various configuration knobs.
135    ///
136    /// [builder-struct]: ./struct.CacheBuilder.html
137    pub fn builder() -> CacheBuilder<K, V, Cache<K, V, RandomState>> {
138        CacheBuilder::default()
139    }
140}
141
142//
143// public
144//
145impl<K, V, S> Cache<K, V, S> {
146    /// Returns a read-only cache policy of this cache.
147    ///
148    /// At this time, cache policy cannot be modified after cache creation.
149    /// A future version may support to modify it.
150    pub fn policy(&self) -> Policy {
151        Policy::new(self.max_capacity)
152    }
153
154    /// Returns the number of entries in this cache.
155    ///
156    /// # Example
157    ///
158    /// ```rust
159    /// use micro_moka::unsync::Cache;
160    ///
161    /// let mut cache = Cache::new(10);
162    /// cache.insert('n', "Netherland Dwarf");
163    /// cache.insert('l', "Lop Eared");
164    /// cache.insert('d', "Dutch");
165    ///
166    /// // Ensure an entry exists.
167    /// assert!(cache.contains_key(&'n'));
168    ///
169    /// // Followings will print the actual numbers.
170    /// println!("{}", cache.entry_count());   // -> 3
171    /// ```
172    ///
173    pub fn entry_count(&self) -> u64 {
174        self.entry_count
175    }
176
177    /// Returns the total weighted size of entries in this cache.
178    ///
179    /// This is equivalent to `entry_count` as weight support has been removed.
180    pub fn weighted_size(&self) -> u64 {
181        self.entry_count
182    }
183}
184
185impl<K, V, S> Cache<K, V, S>
186where
187    K: Hash + Eq,
188    S: BuildHasher + Clone,
189{
190    pub(crate) fn with_everything(
191        max_capacity: Option<u64>,
192        initial_capacity: Option<usize>,
193        build_hasher: S,
194    ) -> Self {
195        let init_cap = initial_capacity.unwrap_or_default();
196
197        Self {
198            max_capacity,
199            entry_count: 0,
200            table: HashTable::with_capacity(init_cap),
201            build_hasher,
202            slab: if init_cap > 0 {
203                Slab::with_capacity(init_cap)
204            } else {
205                Slab::new()
206            },
207            deque: IndexDeque::default(),
208            frequency_sketch: FrequencySketch::default(),
209            frequency_sketch_enabled: false,
210        }
211    }
212
213    /// Returns `true` if the cache contains a value for the key.
214    ///
215    /// Unlike the `get` method, this method is not considered a cache read operation,
216    /// so it does not update the historic popularity estimator.
217    ///
218    /// The key may be any borrowed form of the cache's key type, but `Hash` and `Eq`
219    /// on the borrowed form _must_ match those for the key type.
220    #[inline]
221    pub fn contains_key<Q>(&self, key: &Q) -> bool
222    where
223        K: Borrow<Q>,
224        Q: Hash + Eq + ?Sized,
225    {
226        let hash = self.hash(key);
227        self.table
228            .find(hash, |&idx| self.slab.get(idx).key.borrow() == key)
229            .is_some()
230    }
231
232    /// Returns an immutable reference of the value corresponding to the key.
233    ///
234    /// The key may be any borrowed form of the cache's key type, but `Hash` and `Eq`
235    /// on the borrowed form _must_ match those for the key type.
236    #[inline]
237    pub fn get<Q>(&mut self, key: &Q) -> Option<&V>
238    where
239        K: Borrow<Q>,
240        Q: Hash + Eq + ?Sized,
241    {
242        let hash = self.hash(key);
243        self.frequency_sketch.increment(hash);
244
245        let idx = match self
246            .table
247            .find(hash, |&idx| self.slab.get(idx).key.borrow() == key)
248        {
249            Some(&idx) => idx,
250            None => return None,
251        };
252
253        self.deque.move_to_back(&mut self.slab, idx);
254        Some(&self.slab.get(idx).value)
255    }
256
257    /// Inserts a key-value pair into the cache.
258    ///
259    /// If the cache has this key present, the value is updated.
260    #[inline]
261    pub fn insert(&mut self, key: K, value: V) {
262        let weights_to_evict = self.weights_to_evict();
263        if weights_to_evict > 0 {
264            self.evict_lru_entries(weights_to_evict);
265        }
266
267        let hash = self.hash(&key);
268
269        if let Some(&idx) = self
270            .table
271            .find(hash, |&idx| self.slab.get(idx).key.borrow() == &key)
272        {
273            self.slab.get_mut(idx).value = value;
274            self.deque.move_to_back(&mut self.slab, idx);
275            return;
276        }
277
278        let slab_entry = SlabEntry {
279            key,
280            value,
281            hash,
282            prev: SENTINEL,
283            next: SENTINEL,
284        };
285        let idx = self.slab.allocate(slab_entry);
286
287        let slab = &self.slab;
288        self.table
289            .insert_unique(hash, idx, |&existing_idx| slab.get(existing_idx).hash);
290
291        self.handle_insert(idx, hash);
292    }
293
294    /// Discards any cached value for the key.
295    ///
296    /// The key may be any borrowed form of the cache's key type, but `Hash` and `Eq`
297    /// on the borrowed form _must_ match those for the key type.
298    #[inline]
299    pub fn invalidate<Q>(&mut self, key: &Q)
300    where
301        K: Borrow<Q>,
302        Q: Hash + Eq + ?Sized,
303    {
304        let hash = self.hash(key);
305        let slab = &self.slab;
306        if let Ok(entry) = self
307            .table
308            .find_entry(hash, |&idx| slab.get(idx).key.borrow() == key)
309        {
310            let (idx, _) = entry.remove();
311            self.deque.unlink(&mut self.slab, idx);
312            self.slab.deallocate(idx);
313            self.entry_count -= 1;
314        }
315    }
316
317    /// Discards any cached value for the key, returning the cached value.
318    ///
319    /// The key may be any borrowed form of the cache's key type, but `Hash` and `Eq`
320    /// on the borrowed form _must_ match those for the key type.
321    #[inline]
322    pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
323    where
324        K: Borrow<Q>,
325        Q: Hash + Eq + ?Sized,
326    {
327        let hash = self.hash(key);
328        let slab = &self.slab;
329        if let Ok(entry) = self
330            .table
331            .find_entry(hash, |&idx| slab.get(idx).key.borrow() == key)
332        {
333            let (idx, _) = entry.remove();
334            self.deque.unlink(&mut self.slab, idx);
335            let slab_entry = self.slab.deallocate(idx);
336            self.entry_count -= 1;
337            Some(slab_entry.value)
338        } else {
339            None
340        }
341    }
342
343    /// Discards all cached values.
344    ///
345    /// Like the `invalidate` method, this method does not clear the historic
346    /// popularity estimator of keys so that it retains the client activities of
347    /// trying to retrieve an item.
348    #[cold]
349    #[inline(never)]
350    pub fn invalidate_all(&mut self) {
351        let old_capacity = self.table.capacity();
352        let old_table = std::mem::replace(&mut self.table, HashTable::new());
353        let old_slab = std::mem::replace(&mut self.slab, Slab::new());
354        self.deque.clear();
355        self.entry_count = 0;
356
357        drop(old_table);
358        drop(old_slab);
359
360        self.table.reserve(old_capacity, |&idx| {
361            // This closure is for rehashing during reserve. Since the table is
362            // empty after the swap, this will never be called, but we must
363            // provide it.
364            let _ = idx;
365            0
366        });
367    }
368
369    /// Discards cached values that satisfy a predicate.
370    ///
371    /// `invalidate_entries_if` takes a closure that returns `true` or `false`.
372    /// `invalidate_entries_if` will apply the closure to each cached value,
373    /// and if the closure returns `true`, the value will be invalidated.
374    ///
375    /// Like the `invalidate` method, this method does not clear the historic
376    /// popularity estimator of keys so that it retains the client activities of
377    /// trying to retrieve an item.
378    #[cold]
379    #[inline(never)]
380    pub fn invalidate_entries_if(&mut self, mut predicate: impl FnMut(&K, &V) -> bool) {
381        let indices_to_invalidate: Vec<u32> = self
382            .slab
383            .iter()
384            .filter(|(_, entry)| predicate(&entry.key, &entry.value))
385            .map(|(idx, _)| idx)
386            .collect();
387
388        let mut invalidated = 0u64;
389        for idx in indices_to_invalidate {
390            let hash = self.slab.get(idx).hash;
391            if let Ok(entry) = self.table.find_entry(hash, |&table_idx| table_idx == idx) {
392                entry.remove();
393                self.deque.unlink(&mut self.slab, idx);
394                self.slab.deallocate(idx);
395                invalidated += 1;
396            }
397        }
398        self.entry_count -= invalidated;
399    }
400
401    /// Creates an iterator visiting all key-value pairs in arbitrary order. The
402    /// iterator element type is `(&K, &V)`.
403    ///
404    /// Unlike the `get` method, visiting entries via an iterator do not update the
405    /// historic popularity estimator or reset idle timers for keys.
406    ///
407    /// # Examples
408    ///
409    /// ```rust
410    /// use micro_moka::unsync::Cache;
411    ///
412    /// let mut cache = Cache::new(100);
413    /// cache.insert("Julia", 14);
414    ///
415    /// let mut iter = cache.iter();
416    /// let (k, v) = iter.next().unwrap(); // (&K, &V)
417    /// assert_eq!(k, &"Julia");
418    /// assert_eq!(v, &14);
419    ///
420    /// assert!(iter.next().is_none());
421    /// ```
422    ///
423    pub fn iter(&self) -> Iter<'_, K, V> {
424        Iter::new(&self.slab.entries)
425    }
426}
427
428//
429// private
430//
431impl<K, V, S> Cache<K, V, S>
432where
433    K: Hash + Eq,
434    S: BuildHasher + Clone,
435{
436    #[inline]
437    fn hash<Q>(&self, key: &Q) -> u64
438    where
439        K: Borrow<Q>,
440        Q: Hash + Eq + ?Sized,
441    {
442        self.build_hasher.hash_one(key)
443    }
444
445    #[inline]
446    fn has_enough_capacity(&self, candidate_weight: u32, ws: u64) -> bool {
447        self.max_capacity
448            .map(|limit| ws + candidate_weight as u64 <= limit)
449            .unwrap_or(true)
450    }
451
452    #[inline]
453    fn weights_to_evict(&self) -> u64 {
454        self.max_capacity
455            .map(|limit| self.entry_count.saturating_sub(limit))
456            .unwrap_or_default()
457    }
458
459    #[inline]
460    fn should_enable_frequency_sketch(&self) -> bool {
461        if self.frequency_sketch_enabled {
462            false
463        } else if let Some(max_cap) = self.max_capacity {
464            self.entry_count >= max_cap / 2
465        } else {
466            false
467        }
468    }
469
470    #[inline]
471    fn enable_frequency_sketch(&mut self) {
472        if let Some(max_cap) = self.max_capacity {
473            self.do_enable_frequency_sketch(max_cap);
474        }
475    }
476
477    #[cfg(test)]
478    fn enable_frequency_sketch_for_testing(&mut self) {
479        if let Some(max_cap) = self.max_capacity {
480            self.do_enable_frequency_sketch(max_cap);
481        }
482    }
483
484    #[inline]
485    fn do_enable_frequency_sketch(&mut self, cache_capacity: u64) {
486        let skt_capacity = common::sketch_capacity(cache_capacity);
487        self.frequency_sketch.ensure_capacity(skt_capacity);
488        self.frequency_sketch_enabled = true;
489    }
490
491    #[inline]
492    fn handle_insert(&mut self, idx: u32, hash: u64) {
493        let has_free_space = self.has_enough_capacity(1, self.entry_count);
494
495        if has_free_space {
496            self.deque.push_back(&mut self.slab, idx);
497            self.entry_count += 1;
498
499            if self.should_enable_frequency_sketch() {
500                self.enable_frequency_sketch();
501            }
502            return;
503        }
504
505        if let Some(max) = self.max_capacity {
506            if max == 0 {
507                self.remove_by_index(idx);
508                return;
509            }
510        }
511
512        let candidate_freq = self.frequency_sketch.frequency(hash);
513
514        match self.admit(candidate_freq) {
515            AdmissionResult::Admitted { victim_index } => {
516                self.remove_by_index(victim_index);
517
518                self.deque.push_back(&mut self.slab, idx);
519                self.entry_count += 1;
520
521                if self.should_enable_frequency_sketch() {
522                    self.enable_frequency_sketch();
523                }
524            }
525            AdmissionResult::Rejected => {
526                self.remove_by_index(idx);
527            }
528        }
529    }
530
531    #[inline]
532    fn admit(&self, candidate_freq: u8) -> AdmissionResult {
533        let Some(victim_index) = self.deque.peek_front() else {
534            return AdmissionResult::Rejected;
535        };
536        let victim_hash = self.slab.get(victim_index).hash;
537        let victim_freq = self.frequency_sketch.frequency(victim_hash);
538
539        if candidate_freq > victim_freq {
540            AdmissionResult::Admitted { victim_index }
541        } else {
542            AdmissionResult::Rejected
543        }
544    }
545
546    fn remove_by_index(&mut self, idx: u32) {
547        let hash = self.slab.get(idx).hash;
548        if let Ok(entry) = self.table.find_entry(hash, |&table_idx| table_idx == idx) {
549            entry.remove();
550        }
551        let entry = self.slab.get(idx);
552        if entry.prev != SENTINEL || entry.next != SENTINEL || self.deque.head == idx {
553            self.deque.unlink(&mut self.slab, idx);
554            self.entry_count -= 1;
555        }
556        self.slab.deallocate(idx);
557    }
558
559    #[cold]
560    #[inline(never)]
561    fn evict_lru_entries(&mut self, weights_to_evict: u64) {
562        debug_assert!(weights_to_evict > 0);
563        let mut evicted = 0u64;
564
565        for _ in 0..EVICTION_BATCH_SIZE {
566            if evicted >= weights_to_evict {
567                break;
568            }
569
570            let Some(victim_idx) = self.deque.peek_front() else {
571                break;
572            };
573
574            let victim_hash = self.slab.get(victim_idx).hash;
575            if let Ok(entry) = self
576                .table
577                .find_entry(victim_hash, |&table_idx| table_idx == victim_idx)
578            {
579                entry.remove();
580            }
581
582            self.deque.unlink(&mut self.slab, victim_idx);
583            self.slab.deallocate(victim_idx);
584            evicted += 1;
585        }
586
587        self.entry_count -= evicted;
588    }
589}
590
591#[cfg(test)]
592impl<K, V, S> Cache<K, V, S>
593where
594    K: Hash + Eq,
595    S: BuildHasher + Clone,
596{
597}
598
599enum AdmissionResult {
600    Admitted { victim_index: u32 },
601    Rejected,
602}
603
604#[cfg(test)]
605mod tests {
606    use super::Cache;
607
608    #[test]
609    fn basic_single_thread() {
610        let mut cache = Cache::new(3);
611        cache.enable_frequency_sketch_for_testing();
612
613        cache.insert("a", "alice");
614        cache.insert("b", "bob");
615        assert_eq!(cache.get(&"a"), Some(&"alice"));
616        assert!(cache.contains_key(&"a"));
617        assert!(cache.contains_key(&"b"));
618        assert_eq!(cache.get(&"b"), Some(&"bob"));
619        // counts: a -> 1, b -> 1
620
621        cache.insert("c", "cindy");
622        assert_eq!(cache.get(&"c"), Some(&"cindy"));
623        assert!(cache.contains_key(&"c"));
624        // counts: a -> 1, b -> 1, c -> 1
625
626        assert!(cache.contains_key(&"a"));
627        assert_eq!(cache.get(&"a"), Some(&"alice"));
628        assert_eq!(cache.get(&"b"), Some(&"bob"));
629        assert!(cache.contains_key(&"b"));
630        // counts: a -> 2, b -> 2, c -> 1
631
632        // "d" should not be admitted because its frequency is too low.
633        cache.insert("d", "david"); //   count: d -> 0
634        assert_eq!(cache.get(&"d"), None); //   d -> 1
635        assert!(!cache.contains_key(&"d"));
636
637        cache.insert("d", "david");
638        assert!(!cache.contains_key(&"d"));
639        assert_eq!(cache.get(&"d"), None); //   d -> 2
640
641        // "d" should be admitted and "c" should be evicted
642        // because d's frequency is higher than c's.
643        cache.insert("d", "dennis");
644        assert_eq!(cache.get(&"a"), Some(&"alice"));
645        assert_eq!(cache.get(&"b"), Some(&"bob"));
646        assert_eq!(cache.get(&"c"), None);
647        assert_eq!(cache.get(&"d"), Some(&"dennis"));
648        assert!(cache.contains_key(&"a"));
649        assert!(cache.contains_key(&"b"));
650        assert!(!cache.contains_key(&"c"));
651        assert!(cache.contains_key(&"d"));
652
653        cache.invalidate(&"b");
654        assert_eq!(cache.get(&"b"), None);
655        assert!(!cache.contains_key(&"b"));
656    }
657
658    #[test]
659    fn invalidate_all() {
660        let mut cache = Cache::new(100);
661        cache.enable_frequency_sketch_for_testing();
662
663        cache.insert("a", "alice");
664        cache.insert("b", "bob");
665        cache.insert("c", "cindy");
666        assert_eq!(cache.get(&"a"), Some(&"alice"));
667        assert_eq!(cache.get(&"b"), Some(&"bob"));
668        assert_eq!(cache.get(&"c"), Some(&"cindy"));
669        assert!(cache.contains_key(&"a"));
670        assert!(cache.contains_key(&"b"));
671        assert!(cache.contains_key(&"c"));
672
673        cache.invalidate_all();
674
675        cache.insert("d", "david");
676
677        assert!(cache.get(&"a").is_none());
678        assert!(cache.get(&"b").is_none());
679        assert!(cache.get(&"c").is_none());
680        assert_eq!(cache.get(&"d"), Some(&"david"));
681        assert!(!cache.contains_key(&"a"));
682        assert!(!cache.contains_key(&"b"));
683        assert!(!cache.contains_key(&"c"));
684        assert!(cache.contains_key(&"d"));
685    }
686
687    #[test]
688    fn invalidate_entries_if() {
689        use std::collections::HashSet;
690
691        let mut cache = Cache::new(100);
692        cache.enable_frequency_sketch_for_testing();
693
694        cache.insert(0, "alice");
695        cache.insert(1, "bob");
696        cache.insert(2, "alex");
697
698        assert_eq!(cache.get(&0), Some(&"alice"));
699        assert_eq!(cache.get(&1), Some(&"bob"));
700        assert_eq!(cache.get(&2), Some(&"alex"));
701        assert!(cache.contains_key(&0));
702        assert!(cache.contains_key(&1));
703        assert!(cache.contains_key(&2));
704
705        let names = ["alice", "alex"].iter().cloned().collect::<HashSet<_>>();
706        cache.invalidate_entries_if(move |_k, &v| names.contains(v));
707
708        cache.insert(3, "alice");
709
710        assert!(cache.get(&0).is_none());
711        assert!(cache.get(&2).is_none());
712        assert_eq!(cache.get(&1), Some(&"bob"));
713        // This should survive as it was inserted after calling invalidate_entries_if.
714        assert_eq!(cache.get(&3), Some(&"alice"));
715
716        assert!(!cache.contains_key(&0));
717        assert!(cache.contains_key(&1));
718        assert!(!cache.contains_key(&2));
719        assert!(cache.contains_key(&3));
720
721        assert_eq!(cache.table.len(), 2);
722
723        cache.invalidate_entries_if(|_k, &v| v == "alice");
724        cache.invalidate_entries_if(|_k, &v| v == "bob");
725
726        assert!(cache.get(&1).is_none());
727        assert!(cache.get(&3).is_none());
728
729        assert!(!cache.contains_key(&1));
730        assert!(!cache.contains_key(&3));
731
732        assert_eq!(cache.table.len(), 0);
733    }
734
735    #[cfg_attr(target_pointer_width = "16", ignore)]
736    #[test]
737    fn test_skt_capacity_will_not_overflow() {
738        let pot = |exp| 2u64.pow(exp);
739
740        let ensure_sketch_len = |max_capacity, len, name| {
741            let mut cache = Cache::<u8, u8>::new(max_capacity);
742            cache.enable_frequency_sketch_for_testing();
743            assert_eq!(cache.frequency_sketch.table_len(), len as usize, "{}", name);
744        };
745
746        if cfg!(target_pointer_width = "32") {
747            let pot24 = pot(24);
748            let pot16 = pot(16);
749            ensure_sketch_len(0, 128, "0");
750            ensure_sketch_len(128, 128, "128");
751            ensure_sketch_len(pot16, pot16, "pot16");
752            ensure_sketch_len(pot16 + 1, pot(17), "pot16 + 1");
753            ensure_sketch_len(pot24 - 1, pot24, "pot24 - 1");
754            ensure_sketch_len(pot24, pot24, "pot24");
755            ensure_sketch_len(pot(27), pot24, "pot(27)");
756            ensure_sketch_len(u32::MAX as u64, pot24, "u32::MAX");
757        } else {
758            let pot30 = pot(30);
759            let pot16 = pot(16);
760            ensure_sketch_len(0, 128, "0");
761            ensure_sketch_len(128, 128, "128");
762            ensure_sketch_len(pot16, pot16, "pot16");
763            ensure_sketch_len(pot16 + 1, pot(17), "pot16 + 1");
764
765            if !cfg!(circleci) {
766                ensure_sketch_len(pot30 - 1, pot30, "pot30- 1");
767                ensure_sketch_len(pot30, pot30, "pot30");
768                ensure_sketch_len(u64::MAX, pot30, "u64::MAX");
769            }
770        };
771    }
772
773    #[test]
774    fn remove_decrements_entry_count() {
775        let mut cache = Cache::new(3);
776        cache.insert("a", "alice");
777        cache.insert("b", "bob");
778        assert_eq!(cache.entry_count(), 2);
779
780        let removed = cache.remove(&"a");
781        assert_eq!(removed, Some("alice"));
782        assert_eq!(cache.entry_count(), 1);
783
784        cache.remove(&"nonexistent");
785        assert_eq!(cache.entry_count(), 1);
786
787        cache.remove(&"b");
788        assert_eq!(cache.entry_count(), 0);
789    }
790
791    #[test]
792    fn invalidate_decrements_entry_count() {
793        let mut cache = Cache::new(3);
794        cache.insert("a", "alice");
795        cache.insert("b", "bob");
796        assert_eq!(cache.entry_count(), 2);
797
798        cache.invalidate(&"a");
799        assert_eq!(cache.entry_count(), 1);
800
801        cache.invalidate(&"nonexistent");
802        assert_eq!(cache.entry_count(), 1);
803
804        cache.invalidate(&"b");
805        assert_eq!(cache.entry_count(), 0);
806    }
807
808    #[test]
809    fn insert_after_remove_on_full_cache() {
810        let mut cache = Cache::new(2);
811        cache.insert("a", "alice");
812        cache.insert("b", "bob");
813        assert_eq!(cache.entry_count(), 2);
814
815        cache.remove(&"a");
816        assert_eq!(cache.entry_count(), 1);
817
818        cache.insert("c", "cindy");
819        assert_eq!(cache.entry_count(), 2);
820        assert_eq!(cache.get(&"c"), Some(&"cindy"));
821        assert_eq!(cache.get(&"b"), Some(&"bob"));
822        assert_eq!(cache.get(&"a"), None);
823    }
824
825    #[test]
826    fn insert_after_invalidate_on_full_cache() {
827        let mut cache = Cache::new(2);
828        cache.insert("a", "alice");
829        cache.insert("b", "bob");
830        assert_eq!(cache.entry_count(), 2);
831
832        cache.invalidate(&"a");
833        assert_eq!(cache.entry_count(), 1);
834
835        cache.insert("c", "cindy");
836        assert_eq!(cache.entry_count(), 2);
837        assert_eq!(cache.get(&"c"), Some(&"cindy"));
838        assert_eq!(cache.get(&"b"), Some(&"bob"));
839        assert_eq!(cache.get(&"a"), None);
840    }
841
842    #[test]
843    fn invalidate_all_panic_safety() {
844        use std::panic::catch_unwind;
845        use std::panic::AssertUnwindSafe;
846        use std::sync::atomic::{AtomicU32, Ordering};
847
848        static DROP_COUNT: AtomicU32 = AtomicU32::new(0);
849
850        struct PanicOnDrop {
851            id: u32,
852            should_panic: bool,
853        }
854
855        impl Drop for PanicOnDrop {
856            fn drop(&mut self) {
857                DROP_COUNT.fetch_add(1, Ordering::Relaxed);
858                if self.should_panic {
859                    panic!("intentional panic in drop for id={}", self.id);
860                }
861            }
862        }
863
864        DROP_COUNT.store(0, Ordering::Relaxed);
865        let mut cache = Cache::new(10);
866        cache.insert(
867            1,
868            PanicOnDrop {
869                id: 1,
870                should_panic: false,
871            },
872        );
873        cache.insert(
874            2,
875            PanicOnDrop {
876                id: 2,
877                should_panic: true,
878            },
879        );
880        cache.insert(
881            3,
882            PanicOnDrop {
883                id: 3,
884                should_panic: false,
885            },
886        );
887        assert_eq!(cache.entry_count(), 3);
888
889        let result = catch_unwind(AssertUnwindSafe(|| {
890            cache.invalidate_all();
891        }));
892        assert!(result.is_err());
893
894        assert_eq!(cache.entry_count(), 0);
895        assert_eq!(cache.table.len(), 0);
896
897        cache.insert(
898            4,
899            PanicOnDrop {
900                id: 4,
901                should_panic: false,
902            },
903        );
904        assert_eq!(cache.entry_count(), 1);
905        assert!(cache.contains_key(&4));
906    }
907
908    #[test]
909    fn test_debug_format() {
910        let mut cache = Cache::new(10);
911        cache.insert('a', "alice");
912        cache.insert('b', "bob");
913        cache.insert('c', "cindy");
914
915        let debug_str = format!("{:?}", cache);
916        assert!(debug_str.starts_with('{'));
917        assert!(debug_str.contains(r#"'a': "alice""#));
918        assert!(debug_str.contains(r#"'b': "bob""#));
919        assert!(debug_str.contains(r#"'c': "cindy""#));
920        assert!(debug_str.ends_with('}'));
921    }
922
923    #[test]
924    fn sub_capacity_inserts_skip_eviction() {
925        let mut cache = Cache::new(10);
926        for i in 0u32..5 {
927            cache.insert(i, i * 10);
928        }
929        assert_eq!(cache.entry_count(), 5);
930        for i in 0u32..5 {
931            assert_eq!(cache.get(&i), Some(&(i * 10)));
932        }
933    }
934
935    #[test]
936    fn eviction_triggers_when_over_capacity() {
937        let mut cache = Cache::new(3);
938        cache.enable_frequency_sketch_for_testing();
939
940        cache.insert(1, "a");
941        cache.insert(2, "b");
942        cache.insert(3, "c");
943        assert_eq!(cache.entry_count(), 3);
944
945        for _ in 0..5 {
946            cache.get(&1);
947            cache.get(&2);
948            cache.get(&3);
949        }
950
951        cache.insert(4, "d");
952        assert!(cache.entry_count() <= 3);
953    }
954
955    #[test]
956    fn warmup_to_full_transition() {
957        let mut cache = Cache::new(4);
958        cache.enable_frequency_sketch_for_testing();
959
960        cache.insert(1, "a");
961        cache.insert(2, "b");
962        assert_eq!(cache.entry_count(), 2);
963        assert_eq!(cache.weights_to_evict(), 0);
964
965        cache.insert(3, "c");
966        cache.insert(4, "d");
967        assert_eq!(cache.entry_count(), 4);
968        assert_eq!(cache.weights_to_evict(), 0);
969
970        for _ in 0..5 {
971            cache.get(&1);
972            cache.get(&2);
973            cache.get(&3);
974            cache.get(&4);
975        }
976
977        cache.insert(5, "e");
978        assert!(cache.entry_count() <= 4);
979    }
980
981    #[test]
982    fn invalidate_and_remove_skip_eviction_below_capacity() {
983        let mut cache = Cache::new(10);
984        cache.insert(1, "a");
985        cache.insert(2, "b");
986        cache.insert(3, "c");
987        assert_eq!(cache.entry_count(), 3);
988        assert_eq!(cache.weights_to_evict(), 0);
989
990        cache.invalidate(&1);
991        assert_eq!(cache.entry_count(), 2);
992
993        let val = cache.remove(&2);
994        assert_eq!(val, Some("b"));
995        assert_eq!(cache.entry_count(), 1);
996
997        assert_eq!(cache.get(&3), Some(&"c"));
998    }
999}