diff --git a/src/lfu/tinylfu.rs b/src/lfu/tinylfu.rs index b3ce133..e57208d 100644 --- a/src/lfu/tinylfu.rs +++ b/src/lfu/tinylfu.rs @@ -27,12 +27,12 @@ pub struct TinyLFUBuilder> { marker: PhantomData, } -impl Default for TinyLFUBuilder { +impl + Default> Default for TinyLFUBuilder { fn default() -> Self { Self { samples: 0, size: 0, - key_hasher: Some(DefaultKeyHasher::default()), + key_hasher: Some(KH::default()), false_positive_ratio: Some(DEFAULT_FALSE_POSITIVE_RATIO), marker: Default::default(), } @@ -47,6 +47,17 @@ impl TinyLFUBuilder { } impl> TinyLFUBuilder { + /// The constructor of TinyLFUBuilder + pub fn with_hasher(key_hasher: KH) -> Self { + Self { + samples: 0, + size: 0, + key_hasher: Some(key_hasher), + false_positive_ratio: Some(DEFAULT_FALSE_POSITIVE_RATIO), + marker: Default::default(), + } + } + /// Set the samples of TinyLFU pub fn set_samples(self, samples: usize) -> Self { Self { @@ -387,8 +398,65 @@ impl> TinyLFU { } #[cfg(test)] -mod test { - use crate::lfu::tinylfu::TinyLFU; +pub(crate) mod test { + use core::hash::Hasher; + + use crate::lfu::{tinylfu::TinyLFU, KeyHasher}; + + use super::TinyLFUBuilder; + + #[derive(Default)] + pub(crate) struct PassthroughU64Hasher(u64); + impl core::hash::Hasher for PassthroughU64Hasher { + fn finish(&self) -> u64 { + self.0 + } + + fn write(&mut self, _bytes: &[u8]) { + panic!("Can only write u64"); + } + + fn write_u64(&mut self, i: u64) { + self.0 = i; + } + } + + #[derive(Default)] + pub(crate) struct PassthroughU64KeyHasher {} + + impl KeyHasher for PassthroughU64KeyHasher { + fn hash_key(&self, key: &Q) -> u64 + where + u64: core::borrow::Borrow, + Q: core::hash::Hash + Eq + ?Sized, + { + let mut state = PassthroughU64Hasher::default(); + key.hash(&mut state); + state.finish() + } + } + + #[test] + fn test_custom_hasher() { + let mut l: TinyLFU = TinyLFUBuilder::default() + .set_size(4) + .set_samples(4) + .finalize() + .unwrap(); + assert_eq!(l.hash_key(&0), 0); + assert_eq!(l.hash_key(&10), 10); + + l.increment(&1); + l.increment(&1); + l.increment(&1); + + assert!(l.doorkeeper.contains(1)); + assert_eq!(l.ctr.estimate(1), 2); + + l.increment(&1); + assert!(!l.doorkeeper.contains(1)); + assert_eq!(l.ctr.estimate(1), 1); + } #[test] fn test_increment() { diff --git a/src/lfu/wtinylfu.rs b/src/lfu/wtinylfu.rs index 3cb2e72..b6f8ba5 100644 --- a/src/lfu/wtinylfu.rs +++ b/src/lfu/wtinylfu.rs @@ -36,24 +36,38 @@ pub struct WTinyLFUCacheBuilder< marker: PhantomData, } -impl Default for WTinyLFUCacheBuilder { +impl< + K: Hash + Eq, + KH: KeyHasher + Default, + FH: BuildHasher + Default, + RH: BuildHasher + Default, + WH: BuildHasher + Default, + > Default for WTinyLFUCacheBuilder +{ fn default() -> Self { Self { samples: 0, window_cache_size: 0, main_cache_protected_size: 0, main_cache_probationary_size: 0, - window_cache_hasher: Some(DefaultHashBuilder::default()), - main_cache_protected_hasher: Some(DefaultHashBuilder::default()), - main_cache_probationary_hasher: Some(DefaultHashBuilder::default()), - key_hasher: Some(DefaultKeyHasher::default()), + window_cache_hasher: Some(WH::default()), + main_cache_protected_hasher: Some(FH::default()), + main_cache_probationary_hasher: Some(RH::default()), + key_hasher: Some(KH::default()), false_positive_ratio: Some(DEFAULT_FALSE_POSITIVE_RATIO), marker: Default::default(), } } } -impl WTinyLFUCacheBuilder { +impl< + K: Hash + Eq, + KH: KeyHasher + Default, + FH: BuildHasher + Default, + RH: BuildHasher + Default, + WH: BuildHasher + Default, + > WTinyLFUCacheBuilder +{ /// The constructor of WTinyLFUCacheBuilder pub fn new( window_cache_size: usize, @@ -72,6 +86,27 @@ impl WTinyLFUCacheBuilder { impl, FH: BuildHasher, RH: BuildHasher, WH: BuildHasher> WTinyLFUCacheBuilder { + /// Construct a WTinyLFUCacheBuilder using the custom hashers. + pub fn with_hashers( + key_hasher: KH, + protected_hasher: FH, + probationary_hasher: RH, + window_hasher: WH, + ) -> Self { + Self { + samples: 0, + window_cache_size: 0, + main_cache_protected_size: 0, + main_cache_probationary_size: 0, + window_cache_hasher: Some(window_hasher), + main_cache_protected_hasher: Some(protected_hasher), + main_cache_probationary_hasher: Some(probationary_hasher), + key_hasher: Some(key_hasher), + false_positive_ratio: Some(DEFAULT_FALSE_POSITIVE_RATIO), + marker: Default::default(), + } + } + /// Set the samples of TinyLFU pub fn set_samples(self, samples: usize) -> Self { Self { @@ -400,6 +435,19 @@ impl WTinyLFUCache> { ) .finalize() } +} + +impl, FH: BuildHasher, RH: BuildHasher, WH: BuildHasher> + WTinyLFUCache +{ + /// Creates a WTinyLFUCache according to [`WTinyLFUCacheBuilder`] + /// + /// [`WTinyLFUCacheBuilder`]: struct.WTinyLFUCacheBuilder.html + pub fn from_builder( + builder: WTinyLFUCacheBuilder, + ) -> Result { + builder.finalize() + } /// Returns a [`WTinyLFUCacheBuilder`] with default configurations. /// @@ -429,19 +477,6 @@ impl WTinyLFUCache> { } } -impl, FH: BuildHasher, RH: BuildHasher, WH: BuildHasher> - WTinyLFUCache -{ - /// Creates a WTinyLFUCache according to [`WTinyLFUCacheBuilder`] - /// - /// [`WTinyLFUCacheBuilder`]: struct.WTinyLFUCacheBuilder.html - pub fn from_builder( - builder: WTinyLFUCacheBuilder, - ) -> Result { - builder.finalize() - } -} - impl, FH: BuildHasher, RH: BuildHasher, WH: BuildHasher> Cache for WTinyLFUCache { @@ -654,8 +689,120 @@ impl, FH: BuildHasher, RH: BuildHasher, WH: Bu #[cfg(test)] mod test { + use core::hash::BuildHasher; + + use crate::lfu::tinylfu::test::{PassthroughU64Hasher, PassthroughU64KeyHasher}; use crate::lfu::WTinyLFUCache; - use crate::{Cache, PutResult}; + use crate::{Cache, PutResult, WTinyLFUCacheBuilder}; + + #[derive(Default)] + struct PassthroughBuilderU64Hasher {} + + impl BuildHasher for PassthroughBuilderU64Hasher { + type Hasher = PassthroughU64Hasher; + + fn build_hasher(&self) -> Self::Hasher { + Default::default() + } + } + + #[test] + #[cfg_attr(miri, ignore)] + fn test_wtinylfu_custom_hasher() { + let mut cache: WTinyLFUCache< + u64, + u32, + PassthroughU64KeyHasher, + PassthroughBuilderU64Hasher, + PassthroughBuilderU64Hasher, + PassthroughBuilderU64Hasher, + > = WTinyLFUCacheBuilder::default() + .set_window_cache_size(1) + .set_protected_cache_size(2) + .set_probationary_cache_size(2) + .set_samples(5) + .finalize() + .unwrap(); + assert_eq!(cache.cap(), 5); + assert_eq!(cache.window_cache_cap(), 1); + assert_eq!(cache.main_cache_cap(), 4); + + assert_eq!(cache.put(1, 1), PutResult::Put); + assert!(cache.contains(&1)); + assert_eq!(cache.put(2, 2), PutResult::Put); + assert!(cache.contains(&2)); + assert_eq!(cache.put(3, 3), PutResult::Put); + assert!(cache.contains(&3)); + + // current state + // window cache: (MRU) [3] (LRU) + // probationary cache: (MRU) [2, 1] (LRU) + // protected cache: (MRU) [] (LRU) + assert_eq!(cache.window_cache_len(), 1); + assert_eq!(cache.main_cache_len(), 2); + + assert_eq!(cache.put(2, 22), PutResult::Update(2)); + assert_eq!(cache.put(1, 11), PutResult::Update(1)); + + // current state + // window cache: (MRU) [3] (LRU) + // probationary cache: (MRU) [] (LRU) + // protected cache: (MRU) [1, 2] (LRU) + assert_eq!(cache.slru.peek_lru_from_protected(), Some((&2, &22))); + assert_eq!(cache.slru.peek_mru_from_protected(), Some((&1, &11))); + assert_eq!(cache.window_cache_len(), 1); + assert_eq!(cache.slru.protected_len(), 2); + assert_eq!(cache.slru.probationary_len(), 0); + + assert_eq!(cache.put(3, 33), PutResult::Update(3)); + + // current state + // window cache: (MRU) [2] (LRU) + // probationary cache: (MRU) [] (LRU) + // protected cache: (MRU) [3, 1] (LRU) + assert_eq!(cache.lru.peek_lru(), Some((&2, &22))); + assert_eq!(cache.slru.peek_lru_from_protected(), Some((&1, &11))); + assert_eq!(cache.slru.peek_mru_from_protected(), Some((&3, &33))); + assert_eq!(cache.window_cache_len(), 1); + assert_eq!(cache.slru.protected_len(), 2); + assert_eq!(cache.slru.probationary_len(), 0); + + assert_eq!(cache.put(4, 4), PutResult::Put); + assert_eq!(cache.put(5, 5), PutResult::Put); + + // current state + // window cache: (MRU) [5] (LRU) + // probationary cache: (MRU) [4, 2] (LRU) + // protected cache: (MRU) [3, 1] (LRU) + assert_eq!(cache.lru.peek_lru(), Some((&5, &5))); + assert_eq!(cache.slru.peek_lru_from_probationary(), Some((&2, &22))); + assert_eq!(cache.slru.peek_mru_from_probationary(), Some((&4, &4))); + assert_eq!(cache.lru.len(), 1); + assert_eq!(cache.slru.protected_len(), 2); + assert_eq!(cache.slru.probationary_len(), 2); + + assert_eq!(cache.get(&4), Some(&4)); + assert_eq!(cache.get_mut(&5), Some(&mut 5)); + + // current state + // window cache: (MRU) [5] (LRU) + // probationary cache: (MRU) [1, 2] (LRU) + // protected cache: (MRU) [4, 3] (LRU) + assert_eq!(cache.lru.peek_lru(), Some((&5, &5))); + assert_eq!(cache.slru.peek_lru_from_probationary(), Some((&2, &22))); + assert_eq!(cache.slru.peek_mru_from_probationary(), Some((&1, &11))); + assert_eq!(cache.slru.peek_lru_from_protected(), Some((&3, &33))); + assert_eq!(cache.slru.peek_mru_from_protected(), Some((&4, &4))); + assert_eq!(cache.lru.len(), 1); + assert_eq!(cache.slru.protected_len(), 2); + assert_eq!(cache.slru.probationary_len(), 2); + + assert_eq!(cache.peek(&3), Some(&33)); + assert_eq!(cache.peek_mut(&2), Some(&mut 22)); + + assert_eq!(cache.remove(&3), Some(33)); + assert_eq!(cache.remove(&2), Some(22)); + } #[test] #[cfg_attr(miri, ignore)]