caches/lfu/
wtinylfu.rs

1mod error;
2pub use error::WTinyLFUError;
3
4use crate::lfu::{
5    tinylfu::{TinyLFU, TinyLFUBuilder, TinyLFUError, DEFAULT_FALSE_POSITIVE_RATIO},
6    DefaultKeyHasher, KeyHasher,
7};
8use crate::lru::{SegmentedCache, SegmentedCacheBuilder};
9use crate::{Cache, DefaultHashBuilder, LRUCache, PutResult};
10use core::borrow::Borrow;
11use core::hash::{BuildHasher, Hash};
12use core::marker::PhantomData;
13
14const DEFAULT_WINDOW_CACHE_SIZE_RATIO: f64 = 0.01;
15const DEFAULT_HOT_ITEMS_CACHE_SIZE_RATIO: f64 = 0.80;
16
17/// `WTinyLFUCacheBuilder` is used to help build a [`TinyLFUCache`] with custom configurations.
18///
19/// [`TinyLFUCache`]: struct.TinyLFUCache.html
20pub struct WTinyLFUCacheBuilder<
21    K,
22    KH = DefaultKeyHasher<K>,
23    FH = DefaultHashBuilder,
24    RH = DefaultHashBuilder,
25    WH = DefaultHashBuilder,
26> {
27    samples: usize,
28    window_cache_size: usize,
29    main_cache_protected_size: usize,
30    main_cache_probationary_size: usize,
31    window_cache_hasher: Option<WH>,
32    key_hasher: Option<KH>,
33    main_cache_protected_hasher: Option<FH>,
34    main_cache_probationary_hasher: Option<RH>,
35    false_positive_ratio: Option<f64>,
36    marker: PhantomData<K>,
37}
38
39impl<
40        K: Hash + Eq,
41        KH: KeyHasher<K> + Default,
42        FH: BuildHasher + Default,
43        RH: BuildHasher + Default,
44        WH: BuildHasher + Default,
45    > Default for WTinyLFUCacheBuilder<K, KH, FH, RH, WH>
46{
47    fn default() -> Self {
48        Self {
49            samples: 0,
50            window_cache_size: 0,
51            main_cache_protected_size: 0,
52            main_cache_probationary_size: 0,
53            window_cache_hasher: Some(WH::default()),
54            main_cache_protected_hasher: Some(FH::default()),
55            main_cache_probationary_hasher: Some(RH::default()),
56            key_hasher: Some(KH::default()),
57            false_positive_ratio: Some(DEFAULT_FALSE_POSITIVE_RATIO),
58            marker: Default::default(),
59        }
60    }
61}
62
63impl<
64        K: Hash + Eq,
65        KH: KeyHasher<K> + Default,
66        FH: BuildHasher + Default,
67        RH: BuildHasher + Default,
68        WH: BuildHasher + Default,
69    > WTinyLFUCacheBuilder<K, KH, FH, RH, WH>
70{
71    /// The constructor of WTinyLFUCacheBuilder
72    pub fn new(
73        window_cache_size: usize,
74        protected_cache_size: usize,
75        probationary_cache_size: usize,
76        samples: usize,
77    ) -> Self {
78        Self::default()
79            .set_samples(samples)
80            .set_window_cache_size(window_cache_size)
81            .set_protected_cache_size(protected_cache_size)
82            .set_probationary_cache_size(probationary_cache_size)
83    }
84}
85
86impl<K: Hash + Eq, KH: KeyHasher<K>, FH: BuildHasher, RH: BuildHasher, WH: BuildHasher>
87    WTinyLFUCacheBuilder<K, KH, FH, RH, WH>
88{
89    /// Construct a WTinyLFUCacheBuilder using the custom hashers.
90    pub fn with_hashers(
91        key_hasher: KH,
92        protected_hasher: FH,
93        probationary_hasher: RH,
94        window_hasher: WH,
95    ) -> Self {
96        Self {
97            samples: 0,
98            window_cache_size: 0,
99            main_cache_protected_size: 0,
100            main_cache_probationary_size: 0,
101            window_cache_hasher: Some(window_hasher),
102            main_cache_protected_hasher: Some(protected_hasher),
103            main_cache_probationary_hasher: Some(probationary_hasher),
104            key_hasher: Some(key_hasher),
105            false_positive_ratio: Some(DEFAULT_FALSE_POSITIVE_RATIO),
106            marker: Default::default(),
107        }
108    }
109
110    /// Set the samples of TinyLFU
111    pub fn set_samples(self, samples: usize) -> Self {
112        Self {
113            samples,
114            window_cache_size: self.window_cache_size,
115            main_cache_protected_size: self.main_cache_protected_size,
116            main_cache_probationary_size: self.main_cache_probationary_size,
117            window_cache_hasher: self.window_cache_hasher,
118            key_hasher: self.key_hasher,
119            main_cache_protected_hasher: self.main_cache_protected_hasher,
120            main_cache_probationary_hasher: self.main_cache_probationary_hasher,
121            false_positive_ratio: self.false_positive_ratio,
122            marker: self.marker,
123        }
124    }
125
126    /// Set the window cache size
127    pub fn set_window_cache_size(self, sz: usize) -> Self {
128        Self {
129            samples: self.samples,
130            window_cache_size: sz,
131            main_cache_protected_size: self.main_cache_protected_size,
132            main_cache_probationary_size: self.main_cache_probationary_size,
133            window_cache_hasher: self.window_cache_hasher,
134            main_cache_protected_hasher: self.main_cache_protected_hasher,
135            main_cache_probationary_hasher: self.main_cache_probationary_hasher,
136            key_hasher: self.key_hasher,
137            false_positive_ratio: self.false_positive_ratio,
138            marker: self.marker,
139        }
140    }
141
142    /// Set the protected cache size
143    pub fn set_protected_cache_size(self, sz: usize) -> Self {
144        Self {
145            samples: self.samples,
146            window_cache_size: self.window_cache_size,
147            main_cache_protected_size: sz,
148            main_cache_probationary_size: self.main_cache_probationary_size,
149            window_cache_hasher: self.window_cache_hasher,
150            main_cache_protected_hasher: self.main_cache_protected_hasher,
151            main_cache_probationary_hasher: self.main_cache_probationary_hasher,
152            key_hasher: self.key_hasher,
153            false_positive_ratio: self.false_positive_ratio,
154            marker: self.marker,
155        }
156    }
157
158    /// Set the probationary cache size
159    pub fn set_probationary_cache_size(self, sz: usize) -> Self {
160        Self {
161            samples: self.samples,
162            window_cache_size: self.window_cache_size,
163            main_cache_protected_size: self.main_cache_protected_size,
164            main_cache_probationary_size: sz,
165            window_cache_hasher: self.window_cache_hasher,
166            main_cache_protected_hasher: self.main_cache_protected_hasher,
167            main_cache_probationary_hasher: self.main_cache_probationary_hasher,
168            key_hasher: self.key_hasher,
169            false_positive_ratio: self.false_positive_ratio,
170            marker: self.marker,
171        }
172    }
173
174    /// Set the false positive ratio of TinyLFU
175    pub fn set_false_positive_ratio(self, fpr: f64) -> Self {
176        Self {
177            samples: self.samples,
178            window_cache_size: self.window_cache_size,
179            main_cache_protected_size: self.main_cache_protected_size,
180            main_cache_probationary_size: self.main_cache_probationary_size,
181            window_cache_hasher: self.window_cache_hasher,
182            main_cache_protected_hasher: self.main_cache_protected_hasher,
183            main_cache_probationary_hasher: self.main_cache_probationary_hasher,
184            key_hasher: self.key_hasher,
185            false_positive_ratio: Some(fpr),
186            marker: self.marker,
187        }
188    }
189
190    /// Set the window cache's hash builder
191    pub fn set_window_hasher<NWH: BuildHasher>(
192        self,
193        hasher: NWH,
194    ) -> WTinyLFUCacheBuilder<K, KH, FH, RH, NWH> {
195        WTinyLFUCacheBuilder {
196            samples: self.samples,
197            window_cache_size: self.window_cache_size,
198            main_cache_protected_size: self.main_cache_protected_size,
199            main_cache_probationary_size: self.main_cache_probationary_size,
200            window_cache_hasher: Some(hasher),
201            main_cache_protected_hasher: self.main_cache_protected_hasher,
202            main_cache_probationary_hasher: self.main_cache_probationary_hasher,
203            key_hasher: self.key_hasher,
204            false_positive_ratio: self.false_positive_ratio,
205            marker: self.marker,
206        }
207    }
208
209    /// Set the protected cache's hash builder
210    pub fn set_protected_hasher<NFH: BuildHasher>(
211        self,
212        hasher: NFH,
213    ) -> WTinyLFUCacheBuilder<K, KH, NFH, RH, WH> {
214        WTinyLFUCacheBuilder {
215            samples: self.samples,
216            window_cache_size: self.window_cache_size,
217            main_cache_protected_size: self.main_cache_protected_size,
218            main_cache_probationary_size: self.main_cache_probationary_size,
219            window_cache_hasher: self.window_cache_hasher,
220            main_cache_protected_hasher: Some(hasher),
221            main_cache_probationary_hasher: self.main_cache_probationary_hasher,
222            key_hasher: self.key_hasher,
223            false_positive_ratio: self.false_positive_ratio,
224            marker: self.marker,
225        }
226    }
227
228    /// Set the probationary cache's hash builder
229    pub fn set_probationary_hasher<NRH: BuildHasher>(
230        self,
231        hasher: NRH,
232    ) -> WTinyLFUCacheBuilder<K, KH, FH, NRH, WH> {
233        WTinyLFUCacheBuilder {
234            samples: self.samples,
235            window_cache_size: self.window_cache_size,
236            main_cache_protected_size: self.main_cache_protected_size,
237            main_cache_probationary_size: self.main_cache_probationary_size,
238            window_cache_hasher: self.window_cache_hasher,
239            main_cache_protected_hasher: self.main_cache_protected_hasher,
240            key_hasher: self.key_hasher,
241            main_cache_probationary_hasher: Some(hasher),
242            false_positive_ratio: self.false_positive_ratio,
243            marker: self.marker,
244        }
245    }
246
247    /// Set the key hasher
248    pub fn set_key_hasher<NKH: KeyHasher<K>>(
249        self,
250        hasher: NKH,
251    ) -> WTinyLFUCacheBuilder<K, NKH, FH, RH, WH> {
252        WTinyLFUCacheBuilder {
253            samples: self.samples,
254            window_cache_size: self.window_cache_size,
255            main_cache_protected_size: self.main_cache_protected_size,
256            main_cache_probationary_size: self.main_cache_probationary_size,
257            window_cache_hasher: self.window_cache_hasher,
258            main_cache_protected_hasher: self.main_cache_protected_hasher,
259            key_hasher: Some(hasher),
260            main_cache_probationary_hasher: self.main_cache_probationary_hasher,
261            false_positive_ratio: self.false_positive_ratio,
262            marker: self.marker,
263        }
264    }
265
266    /// Finalize the builder to [`TinyLFUCache`]
267    ///
268    /// [`TinyLFUCache`]: struct.TinyLFUCache.html
269    pub fn finalize<V>(self) -> Result<WTinyLFUCache<K, V, KH, FH, RH, WH>, WTinyLFUError>
270    where
271        K: Eq,
272    {
273        if self.window_cache_size == 0 {
274            return Err(WTinyLFUError::InvalidWindowCacheSize(0));
275        }
276
277        if self.main_cache_protected_size == 0 {
278            return Err(WTinyLFUError::InvalidProtectedCacheSize(0));
279        }
280
281        if self.main_cache_probationary_size == 0 {
282            return Err(WTinyLFUError::InvalidProbationaryCacheSize(0));
283        }
284
285        if self.samples == 0 {
286            return Err(WTinyLFUError::InvalidSamples(self.samples));
287        }
288
289        let fp_ratio = self.false_positive_ratio.unwrap();
290        if fp_ratio <= 0.0 || fp_ratio >= 1.0 {
291            return Err(WTinyLFUError::InvalidFalsePositiveRatio(fp_ratio));
292        }
293
294        let lru = LRUCache::with_hasher(self.window_cache_size, self.window_cache_hasher.unwrap())
295            .unwrap();
296
297        let slru = SegmentedCacheBuilder::new(
298            self.main_cache_probationary_size,
299            self.main_cache_protected_size,
300        )
301        .set_probationary_hasher(self.main_cache_probationary_hasher.unwrap())
302        .set_protected_hasher(self.main_cache_protected_hasher.unwrap())
303        .finalize()
304        .unwrap();
305
306        let size = self.window_cache_size
307            + self.main_cache_protected_size
308            + self.main_cache_probationary_size;
309
310        let tinylfu = TinyLFUBuilder::new(size, self.samples)
311            .set_key_hasher(self.key_hasher.unwrap())
312            .set_false_positive_ratio(fp_ratio)
313            .finalize()
314            .map_err(|e| match e {
315                TinyLFUError::InvalidCountMinWidth(v) => WTinyLFUError::InvalidCountMinWidth(v),
316                TinyLFUError::InvalidSamples(v) => WTinyLFUError::InvalidSamples(v),
317                TinyLFUError::InvalidFalsePositiveRatio(v) => {
318                    WTinyLFUError::InvalidFalsePositiveRatio(v)
319                }
320            })?;
321
322        Ok(WTinyLFUCache { tinylfu, lru, slru })
323    }
324}
325
326/// WTinyLFUCache implements the W-TinyLFU, based on the paper [TinyLFU: A Highly Efficient Cache Admission Policy]
327///
328///
329/// # Example
330/// ```rust
331/// use caches::{WTinyLFUCache, PutResult, Cache};
332///
333/// let mut cache = WTinyLFUCache::with_sizes(1, 2, 2, 5).unwrap();
334/// assert_eq!(cache.cap(), 5);
335/// assert_eq!(cache.window_cache_cap(), 1);
336/// assert_eq!(cache.main_cache_cap(), 4);
337///
338/// assert_eq!(cache.put(1, 1), PutResult::Put);
339/// assert!(cache.contains(&1));
340/// assert_eq!(cache.put(2, 2), PutResult::Put);
341/// assert!(cache.contains(&2));
342/// assert_eq!(cache.put(3, 3), PutResult::Put);
343/// assert!(cache.contains(&3));
344///
345/// // current state
346/// // window cache:        (MRU) [3] (LRU)
347/// // probationary cache:  (MRU) [2, 1] (LRU)
348/// // protected cache:     (MRU) [] (LRU)
349/// assert_eq!(cache.window_cache_len(), 1);
350/// assert_eq!(cache.main_cache_len(), 2);
351///
352/// assert_eq!(cache.put(2, 22), PutResult::Update(2));
353/// assert_eq!(cache.put(1, 11), PutResult::Update(1));
354///
355/// // current state
356/// // window cache:        (MRU) [3] (LRU)
357/// // probationary cache:  (MRU) [] (LRU)
358/// // protected cache:     (MRU) [1, 2] (LRU)
359/// assert_eq!(cache.window_cache_len(), 1);
360///
361/// assert_eq!(cache.put(3, 33), PutResult::Update(3));
362///
363/// // current state
364/// // window cache:        (MRU) [2] (LRU)
365/// // probationary cache:  (MRU) [] (LRU)
366/// // protected cache:     (MRU) [3, 1] (LRU)
367/// assert_eq!(cache.window_cache_len(), 1);
368///
369/// assert_eq!(cache.put(4, 4), PutResult::Put);
370/// assert_eq!(cache.put(5, 5), PutResult::Put);
371///
372/// // current state
373/// // window cache:        (MRU) [5] (LRU)
374/// // probationary cache:  (MRU) [4, 2] (LRU)
375/// // protected cache:     (MRU) [3, 1] (LRU)
376///
377/// assert_eq!(cache.get(&4), Some(&4));
378/// assert_eq!(cache.get_mut(&5), Some(&mut 5));
379///
380/// // current state
381/// // window cache:        (MRU) [5] (LRU)
382/// // probationary cache:  (MRU) [1, 2] (LRU)
383/// // protected cache:     (MRU) [4, 3] (LRU)
384/// assert_eq!(cache.peek(&3), Some(&33));
385/// assert_eq!(cache.peek_mut(&2), Some(&mut 22));
386///
387/// assert_eq!(cache.remove(&3), Some(33));
388/// assert_eq!(cache.remove(&2), Some(22));
389/// ```
390///
391/// [TinyLFU: A Highly Efficient Cache Admission Policy]: https://arxiv.org/pdf/1512.00727.pdf
392pub struct WTinyLFUCache<
393    K: Hash,
394    V,
395    KH = DefaultKeyHasher<K>,
396    FH = DefaultHashBuilder,
397    RH = DefaultHashBuilder,
398    WH = DefaultHashBuilder,
399> {
400    tinylfu: TinyLFU<K, KH>,
401    lru: LRUCache<K, V, WH>,
402    slru: SegmentedCache<K, V, FH, RH>,
403}
404
405impl<K: Hash + Eq, V> WTinyLFUCache<K, V, DefaultKeyHasher<K>> {
406    /// Returns a WTinyLFUCache based on the size and samples
407    ///
408    /// **NOTE:** the size is not the actual cache size,
409    /// the actual size is calculated base on the size param.
410    pub fn new(size: usize, samples: usize) -> Result<Self, WTinyLFUError> {
411        let wsz = ((size as f64) * DEFAULT_WINDOW_CACHE_SIZE_RATIO) as usize;
412        let hsz = ((size as f64) * DEFAULT_HOT_ITEMS_CACHE_SIZE_RATIO) as usize;
413        let csz = ((size as f64) * (1f64 - DEFAULT_HOT_ITEMS_CACHE_SIZE_RATIO)) as usize;
414
415        WTinyLFUCacheBuilder::new(wsz, hsz, csz, samples).finalize()
416    }
417
418    /// Returns a WTinyLFUCache based on the related cache sizes and samples
419    ///
420    /// # Note
421    /// According to the [TinyLFU: A Highly Efficient Cache Admission Policy]
422    ///
423    /// [TinyLFU: A Highly Efficient Cache Admission Policy]: https://arxiv.org/pdf/1512.00727.pdf
424    pub fn with_sizes(
425        window_cache_size: usize,
426        protected_cache_size: usize,
427        probationary_cache_size: usize,
428        samples: usize,
429    ) -> Result<Self, WTinyLFUError> {
430        WTinyLFUCacheBuilder::new(
431            window_cache_size,
432            protected_cache_size,
433            probationary_cache_size,
434            samples,
435        )
436        .finalize()
437    }
438}
439
440impl<K: Hash + Eq, V, KH: KeyHasher<K>, FH: BuildHasher, RH: BuildHasher, WH: BuildHasher>
441    WTinyLFUCache<K, V, KH, FH, RH, WH>
442{
443    /// Creates a WTinyLFUCache according to [`WTinyLFUCacheBuilder`]
444    ///
445    /// [`WTinyLFUCacheBuilder`]: struct.WTinyLFUCacheBuilder.html
446    pub fn from_builder(
447        builder: WTinyLFUCacheBuilder<K, KH, FH, RH, WH>,
448    ) -> Result<Self, WTinyLFUError> {
449        builder.finalize()
450    }
451
452    /// Returns a [`WTinyLFUCacheBuilder`] with default configurations.
453    ///
454    /// [`WTinyLFUCacheBuilder`]: struct.WTinyLFUCacheBuilder.html
455    pub fn builder() -> WTinyLFUCacheBuilder<K> {
456        WTinyLFUCacheBuilder::default()
457    }
458
459    /// Returns the window cache len
460    pub fn window_cache_len(&self) -> usize {
461        self.lru.len()
462    }
463
464    /// Returns the window cache cap
465    pub fn window_cache_cap(&self) -> usize {
466        self.lru.cap()
467    }
468
469    /// Returns the main cache len
470    pub fn main_cache_len(&self) -> usize {
471        self.slru.len()
472    }
473
474    /// Returns the main cache cap
475    pub fn main_cache_cap(&self) -> usize {
476        self.slru.cap()
477    }
478}
479
480impl<K: Hash + Eq, V, KH: KeyHasher<K>, FH: BuildHasher, RH: BuildHasher, WH: BuildHasher>
481    Cache<K, V> for WTinyLFUCache<K, V, KH, FH, RH, WH>
482{
483    /// Puts a key-value pair into cache, returns a [`PutResult`].
484    ///
485    /// # Example
486    ///
487    /// ```
488    /// use caches::{Cache, WTinyLFUCache};
489    /// use caches::PutResult;
490    /// let mut cache = WTinyLFUCache::with_sizes(1, 2, 2, 5).unwrap();
491    ///
492    /// assert_eq!(PutResult::Put, cache.put(1, "a"));
493    /// assert_eq!(PutResult::Put, cache.put(2, "b"));
494    /// assert_eq!(PutResult::Update("b"), cache.put(2, "beta"));
495    /// assert_eq!(PutResult::Put, cache.put(3, "c"));
496    ///
497    /// assert_eq!(cache.get(&1), Some(&"a"));
498    /// assert_eq!(cache.get(&2), Some(&"beta"));
499    /// ```
500    ///
501    /// [`PutResult`]: struct.PutResult.html
502    fn put(&mut self, k: K, v: V) -> PutResult<K, V> {
503        match self.lru.remove(&k) {
504            None => {
505                if self.slru.contains(&k) {
506                    return self.slru.put(k, v);
507                }
508
509                match self.lru.put(k, v) {
510                    PutResult::Put => PutResult::Put,
511                    PutResult::Update(v) => PutResult::Update(v),
512                    PutResult::Evicted { key, value } => {
513                        if self.slru.len() < self.slru.cap() {
514                            return self.slru.put(key, value);
515                        }
516
517                        match self.slru.peek_lru_from_probationary() {
518                            None => self.slru.put(key, value),
519                            Some((lruk, _)) => {
520                                if self.tinylfu.lt(&key, lruk) {
521                                    PutResult::Evicted { key, value }
522                                } else {
523                                    self.slru.put(key, value)
524                                }
525                            }
526                        }
527                    }
528                    // this branch will never happen,
529                    // we keep this for good measure.
530                    _ => PutResult::Put,
531                }
532            }
533            Some(old) => {
534                if self.slru.protected_len() >= self.slru.protected_cap() {
535                    let ent = self.slru.remove_lru_from_protected().unwrap();
536                    self.lru.put(ent.0, ent.1);
537                }
538
539                self.slru.put_protected(k, v);
540                PutResult::Update(old)
541            }
542        }
543    }
544
545    /// Returns a reference to the value of the key in the cache or `None`.
546    ///
547    /// # Example
548    ///
549    /// ```
550    /// use caches::{Cache, WTinyLFUCache};
551    /// let mut cache = WTinyLFUCache::with_sizes(1, 2, 2, 5).unwrap();
552    ///
553    /// cache.put("apple", 8);
554    /// cache.put("banana", 4);
555    /// cache.put("banana", 6);
556    /// cache.put("pear", 2);
557    ///
558    /// assert_eq!(cache.get(&"banana"), Some(&6));
559    /// ```
560    fn get<Q>(&mut self, k: &Q) -> Option<&V>
561    where
562        K: Borrow<Q>,
563        Q: Hash + Eq + ?Sized,
564    {
565        self.tinylfu.try_reset();
566        self.tinylfu.increment(k);
567
568        match self.lru.get(k) {
569            Some(v) => Some(v),
570            None => self.slru.get(k),
571        }
572    }
573
574    /// Returns a mutable reference to the value of the key in the cache or `None`.
575    ///
576    /// # Example
577    ///
578    /// ```
579    /// use caches::{Cache, WTinyLFUCache};
580    /// let mut cache = WTinyLFUCache::with_sizes(1, 2, 2, 5).unwrap();
581    ///
582    /// cache.put("apple", 8);
583    /// cache.put("banana", 4);
584    /// cache.put("banana", 6);
585    /// cache.put("pear", 2);
586    ///
587    /// assert_eq!(cache.get_mut(&"banana"), Some(&mut 6));
588    /// ```
589    fn get_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
590    where
591        K: Borrow<Q>,
592        Q: Hash + Eq + ?Sized,
593    {
594        self.tinylfu.try_reset();
595        self.tinylfu.increment(k);
596        match self.lru.get_mut(k) {
597            Some(v) => Some(v),
598            None => self.slru.get_mut(k),
599        }
600    }
601
602    /// Returns a reference to the value corresponding to the key in the cache or `None` if it is
603    /// not present in the cache. Unlike `get`, `peek` does not update the LRU list so the key's
604    /// position will be unchanged.
605    ///
606    /// # Example
607    ///
608    /// ```
609    /// use caches::{Cache, WTinyLFUCache};
610    ///
611    /// let mut cache = WTinyLFUCache::with_sizes(1, 2, 2, 5).unwrap();
612    ///
613    /// cache.put(1, "a");
614    /// cache.put(2, "b");
615    ///
616    /// assert_eq!(cache.peek(&1), Some(&"a"));
617    /// assert_eq!(cache.peek(&2), Some(&"b"));
618    /// ```
619    fn peek<Q>(&self, k: &Q) -> Option<&V>
620    where
621        K: Borrow<Q>,
622        Q: Hash + Eq + ?Sized,
623    {
624        self.lru.peek(k).or_else(|| self.slru.peek(k))
625    }
626
627    /// Returns a mutable reference to the value corresponding to the key in the cache or `None`
628    /// if it is not present in the cache. Unlike `get_mut`, `peek_mut` does not update the LRU
629    /// list so the key's position will be unchanged.
630    ///
631    /// # Example
632    ///
633    /// ```
634    /// use caches::{Cache, WTinyLFUCache};
635    ///
636    /// let mut cache = WTinyLFUCache::with_sizes(1, 2, 2, 5).unwrap();
637    ///
638    /// cache.put(1, "a");
639    /// cache.put(2, "b");
640    ///
641    /// assert_eq!(cache.peek_mut(&1), Some(&mut "a"));
642    /// assert_eq!(cache.peek_mut(&2), Some(&mut "b"));
643    /// ```
644    fn peek_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
645    where
646        K: Borrow<Q>,
647        Q: Hash + Eq + ?Sized,
648    {
649        match self.lru.peek_mut(k) {
650            Some(v) => Some(v),
651            None => self.slru.peek_mut(k),
652        }
653    }
654
655    fn contains<Q>(&self, k: &Q) -> bool
656    where
657        K: Borrow<Q>,
658        Q: Eq + Hash + ?Sized,
659    {
660        self.lru.contains(k) || self.slru.contains(k)
661    }
662
663    fn remove<Q>(&mut self, k: &Q) -> Option<V>
664    where
665        K: Borrow<Q>,
666        Q: Eq + Hash + ?Sized,
667    {
668        self.lru.remove(k).or_else(|| self.slru.remove(k))
669    }
670
671    fn purge(&mut self) {
672        self.lru.purge();
673        self.slru.purge();
674        self.tinylfu.clear();
675    }
676
677    fn len(&self) -> usize {
678        self.lru.len() + self.slru.len()
679    }
680
681    fn cap(&self) -> usize {
682        self.lru.cap() + self.slru.cap()
683    }
684
685    fn is_empty(&self) -> bool {
686        self.lru.is_empty() && self.slru.is_empty()
687    }
688}
689
690impl<
691        K: Hash + Eq + Clone,
692        V: Clone,
693        KH: KeyHasher<K> + Clone,
694        FH: BuildHasher + Clone,
695        RH: BuildHasher + Clone,
696        WH: BuildHasher + Clone,
697    > Clone for WTinyLFUCache<K, V, KH, FH, RH, WH>
698{
699    fn clone(&self) -> Self {
700        Self {
701            tinylfu: self.tinylfu.clone(),
702            lru: self.lru.clone(),
703            slru: self.slru.clone(),
704        }
705    }
706}
707
708#[cfg(test)]
709mod test {
710    use core::hash::BuildHasher;
711
712    use crate::lfu::tinylfu::test::{PassthroughU64Hasher, PassthroughU64KeyHasher};
713    use crate::lfu::WTinyLFUCache;
714    use crate::{Cache, PutResult, WTinyLFUCacheBuilder};
715
716    #[derive(Default)]
717    struct PassthroughBuilderU64Hasher {}
718
719    impl BuildHasher for PassthroughBuilderU64Hasher {
720        type Hasher = PassthroughU64Hasher;
721
722        fn build_hasher(&self) -> Self::Hasher {
723            Default::default()
724        }
725    }
726
727    #[test]
728    fn test_wtinylfu_custom_hasher() {
729        let mut cache: WTinyLFUCache<
730            u64,
731            u32,
732            PassthroughU64KeyHasher,
733            PassthroughBuilderU64Hasher,
734            PassthroughBuilderU64Hasher,
735            PassthroughBuilderU64Hasher,
736        > = WTinyLFUCacheBuilder::default()
737            .set_window_cache_size(1)
738            .set_protected_cache_size(2)
739            .set_probationary_cache_size(2)
740            .set_samples(5)
741            .finalize()
742            .unwrap();
743        assert_eq!(cache.cap(), 5);
744        assert_eq!(cache.window_cache_cap(), 1);
745        assert_eq!(cache.main_cache_cap(), 4);
746
747        assert_eq!(cache.put(1, 1), PutResult::Put);
748        assert!(cache.contains(&1));
749        assert_eq!(cache.put(2, 2), PutResult::Put);
750        assert!(cache.contains(&2));
751        assert_eq!(cache.put(3, 3), PutResult::Put);
752        assert!(cache.contains(&3));
753
754        // current state
755        // window cache:        (MRU) [3] (LRU)
756        // probationary cache:  (MRU) [2, 1] (LRU)
757        // protected cache:     (MRU) [] (LRU)
758        assert_eq!(cache.window_cache_len(), 1);
759        assert_eq!(cache.main_cache_len(), 2);
760
761        assert_eq!(cache.put(2, 22), PutResult::Update(2));
762        assert_eq!(cache.put(1, 11), PutResult::Update(1));
763
764        // current state
765        // window cache:        (MRU) [3] (LRU)
766        // probationary cache:  (MRU) [] (LRU)
767        // protected cache:     (MRU) [1, 2] (LRU)
768        assert_eq!(cache.slru.peek_lru_from_protected(), Some((&2, &22)));
769        assert_eq!(cache.slru.peek_mru_from_protected(), Some((&1, &11)));
770        assert_eq!(cache.window_cache_len(), 1);
771        assert_eq!(cache.slru.protected_len(), 2);
772        assert_eq!(cache.slru.probationary_len(), 0);
773
774        assert_eq!(cache.put(3, 33), PutResult::Update(3));
775
776        // current state
777        // window cache:        (MRU) [2] (LRU)
778        // probationary cache:  (MRU) [] (LRU)
779        // protected cache:     (MRU) [3, 1] (LRU)
780        assert_eq!(cache.lru.peek_lru(), Some((&2, &22)));
781        assert_eq!(cache.slru.peek_lru_from_protected(), Some((&1, &11)));
782        assert_eq!(cache.slru.peek_mru_from_protected(), Some((&3, &33)));
783        assert_eq!(cache.window_cache_len(), 1);
784        assert_eq!(cache.slru.protected_len(), 2);
785        assert_eq!(cache.slru.probationary_len(), 0);
786
787        assert_eq!(cache.put(4, 4), PutResult::Put);
788        assert_eq!(cache.put(5, 5), PutResult::Put);
789
790        // current state
791        // window cache:        (MRU) [5] (LRU)
792        // probationary cache:  (MRU) [4, 2] (LRU)
793        // protected cache:     (MRU) [3, 1] (LRU)
794        assert_eq!(cache.lru.peek_lru(), Some((&5, &5)));
795        assert_eq!(cache.slru.peek_lru_from_probationary(), Some((&2, &22)));
796        assert_eq!(cache.slru.peek_mru_from_probationary(), Some((&4, &4)));
797        assert_eq!(cache.lru.len(), 1);
798        assert_eq!(cache.slru.protected_len(), 2);
799        assert_eq!(cache.slru.probationary_len(), 2);
800
801        assert_eq!(cache.get(&4), Some(&4));
802        assert_eq!(cache.get_mut(&5), Some(&mut 5));
803
804        // current state
805        // window cache:        (MRU) [5] (LRU)
806        // probationary cache:  (MRU) [1, 2] (LRU)
807        // protected cache:     (MRU) [4, 3] (LRU)
808        assert_eq!(cache.lru.peek_lru(), Some((&5, &5)));
809        assert_eq!(cache.slru.peek_lru_from_probationary(), Some((&2, &22)));
810        assert_eq!(cache.slru.peek_mru_from_probationary(), Some((&1, &11)));
811        assert_eq!(cache.slru.peek_lru_from_protected(), Some((&3, &33)));
812        assert_eq!(cache.slru.peek_mru_from_protected(), Some((&4, &4)));
813        assert_eq!(cache.lru.len(), 1);
814        assert_eq!(cache.slru.protected_len(), 2);
815        assert_eq!(cache.slru.probationary_len(), 2);
816
817        assert_eq!(cache.peek(&3), Some(&33));
818        assert_eq!(cache.peek_mut(&2), Some(&mut 22));
819
820        assert_eq!(cache.remove(&3), Some(33));
821        assert_eq!(cache.remove(&2), Some(22));
822    }
823
824    #[test]
825    fn test_wtinylfu() {
826        let mut cache = WTinyLFUCache::with_sizes(1, 2, 2, 5).unwrap();
827        assert_eq!(cache.cap(), 5);
828        assert_eq!(cache.window_cache_cap(), 1);
829        assert_eq!(cache.main_cache_cap(), 4);
830
831        assert_eq!(cache.put(1, 1), PutResult::Put);
832        assert!(cache.contains(&1));
833        assert_eq!(cache.put(2, 2), PutResult::Put);
834        assert!(cache.contains(&2));
835        assert_eq!(cache.put(3, 3), PutResult::Put);
836        assert!(cache.contains(&3));
837
838        // current state
839        // window cache:        (MRU) [3] (LRU)
840        // probationary cache:  (MRU) [2, 1] (LRU)
841        // protected cache:     (MRU) [] (LRU)
842        assert_eq!(cache.window_cache_len(), 1);
843        assert_eq!(cache.main_cache_len(), 2);
844
845        assert_eq!(cache.put(2, 22), PutResult::Update(2));
846        assert_eq!(cache.put(1, 11), PutResult::Update(1));
847
848        // current state
849        // window cache:        (MRU) [3] (LRU)
850        // probationary cache:  (MRU) [] (LRU)
851        // protected cache:     (MRU) [1, 2] (LRU)
852        assert_eq!(cache.slru.peek_lru_from_protected(), Some((&2, &22)));
853        assert_eq!(cache.slru.peek_mru_from_protected(), Some((&1, &11)));
854        assert_eq!(cache.window_cache_len(), 1);
855        assert_eq!(cache.slru.protected_len(), 2);
856        assert_eq!(cache.slru.probationary_len(), 0);
857
858        assert_eq!(cache.put(3, 33), PutResult::Update(3));
859
860        // current state
861        // window cache:        (MRU) [2] (LRU)
862        // probationary cache:  (MRU) [] (LRU)
863        // protected cache:     (MRU) [3, 1] (LRU)
864        assert_eq!(cache.lru.peek_lru(), Some((&2, &22)));
865        assert_eq!(cache.slru.peek_lru_from_protected(), Some((&1, &11)));
866        assert_eq!(cache.slru.peek_mru_from_protected(), Some((&3, &33)));
867        assert_eq!(cache.window_cache_len(), 1);
868        assert_eq!(cache.slru.protected_len(), 2);
869        assert_eq!(cache.slru.probationary_len(), 0);
870
871        assert_eq!(cache.put(4, 4), PutResult::Put);
872        assert_eq!(cache.put(5, 5), PutResult::Put);
873
874        // current state
875        // window cache:        (MRU) [5] (LRU)
876        // probationary cache:  (MRU) [4, 2] (LRU)
877        // protected cache:     (MRU) [3, 1] (LRU)
878        assert_eq!(cache.lru.peek_lru(), Some((&5, &5)));
879        assert_eq!(cache.slru.peek_lru_from_probationary(), Some((&2, &22)));
880        assert_eq!(cache.slru.peek_mru_from_probationary(), Some((&4, &4)));
881        assert_eq!(cache.lru.len(), 1);
882        assert_eq!(cache.slru.protected_len(), 2);
883        assert_eq!(cache.slru.probationary_len(), 2);
884
885        assert_eq!(cache.get(&4), Some(&4));
886        assert_eq!(cache.get_mut(&5), Some(&mut 5));
887
888        // current state
889        // window cache:        (MRU) [5] (LRU)
890        // probationary cache:  (MRU) [1, 2] (LRU)
891        // protected cache:     (MRU) [4, 3] (LRU)
892        assert_eq!(cache.lru.peek_lru(), Some((&5, &5)));
893        assert_eq!(cache.slru.peek_lru_from_probationary(), Some((&2, &22)));
894        assert_eq!(cache.slru.peek_mru_from_probationary(), Some((&1, &11)));
895        assert_eq!(cache.slru.peek_lru_from_protected(), Some((&3, &33)));
896        assert_eq!(cache.slru.peek_mru_from_protected(), Some((&4, &4)));
897        assert_eq!(cache.lru.len(), 1);
898        assert_eq!(cache.slru.protected_len(), 2);
899        assert_eq!(cache.slru.probationary_len(), 2);
900
901        assert_eq!(cache.peek(&3), Some(&33));
902        assert_eq!(cache.peek_mut(&2), Some(&mut 22));
903
904        assert_eq!(cache.remove(&3), Some(33));
905        assert_eq!(cache.remove(&2), Some(22));
906    }
907
908    #[test]
909    fn test_wtinylfu_clone() {
910        let mut cache = WTinyLFUCache::with_sizes(1, 2, 2, 5).unwrap();
911        assert_eq!(cache.cap(), 5);
912        assert_eq!(cache.window_cache_cap(), 1);
913        assert_eq!(cache.main_cache_cap(), 4);
914
915        assert_eq!(cache.put(1, 1), PutResult::Put);
916        assert!(cache.contains(&1));
917        assert_eq!(cache.put(2, 2), PutResult::Put);
918        assert!(cache.contains(&2));
919        assert_eq!(cache.put(3, 3), PutResult::Put);
920        assert!(cache.contains(&3));
921
922        assert_eq!(cache.put(4, 3), PutResult::Evicted { key: 1, value: 1 });
923    }
924}
caches/lfu/wtinylfu.rs

caches/lfu/
wtinylfu.rs
