Make custom hashers actually usable

Previously the functions for TinyLFU to construct with custom hashes were unavailable, making it impossible to actually construct the caches with custom hashers. Rearrange the code so that constructors are available for custom hashers when using the builder.
al8n · Oct 30, 2023 · 1b92eb5 · 1b92eb5
1 parent 7f23283
commit 1b92eb5
Show file tree

Hide file tree

Showing 2 changed files with 208 additions and 25 deletions.
diff --git a/src/lfu/tinylfu.rs b/src/lfu/tinylfu.rs
@@ -27,12 +27,12 @@ pub struct TinyLFUBuilder<K, KH = DefaultKeyHasher<K>> {
     marker: PhantomData<K>,
 }
 
-impl<K: Hash + Eq> Default for TinyLFUBuilder<K> {
+impl<K: Hash + Eq, KH: KeyHasher<K> + Default> Default for TinyLFUBuilder<K, KH> {
     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<K: Hash + Eq> TinyLFUBuilder<K> {
 }
 
 impl<K: Hash + Eq, KH: KeyHasher<K>> TinyLFUBuilder<K, KH> {
+    /// 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,60 @@ impl<K: Hash + Eq, KH: KeyHasher<K>> TinyLFU<K, KH> {
 }
 
 #[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<u64> for PassthroughU64KeyHasher {
+        fn hash_key<Q>(&self, key: &Q) -> u64
+        where
+            u64: core::borrow::Borrow<Q>,
+            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<u64, PassthroughU64KeyHasher> = 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
@@ -7,7 +7,7 @@ use crate::lfu::{
 };
 use crate::lru::{SegmentedCache, SegmentedCacheBuilder};
 use crate::{Cache, DefaultHashBuilder, LRUCache, PutResult};
-use core::borrow::Borrow;
+use core::{borrow::Borrow};
 use core::hash::{BuildHasher, Hash};
 use core::marker::PhantomData;
 
@@ -36,24 +36,24 @@ pub struct WTinyLFUCacheBuilder<
     marker: PhantomData<K>,
 }
 
-impl<K: Hash + Eq> Default for WTinyLFUCacheBuilder<K> {
+impl<K: Hash + Eq, KH: KeyHasher<K> + Default, FH: BuildHasher + Default, RH: BuildHasher + Default, WH: BuildHasher + Default> Default for WTinyLFUCacheBuilder<K, KH, FH, RH, WH> {
     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<K: Hash + Eq> WTinyLFUCacheBuilder<K> {
+impl<K: Hash + Eq, KH: KeyHasher<K> + Default, FH: BuildHasher + Default, RH: BuildHasher + Default, WH: BuildHasher + Default> WTinyLFUCacheBuilder<K, KH, FH, RH, WH> {
     /// The constructor of WTinyLFUCacheBuilder
     pub fn new(
         window_cache_size: usize,
@@ -72,6 +72,27 @@ impl<K: Hash + Eq> WTinyLFUCacheBuilder<K> {
 impl<K: Hash + Eq, KH: KeyHasher<K>, FH: BuildHasher, RH: BuildHasher, WH: BuildHasher>
     WTinyLFUCacheBuilder<K, KH, FH, RH, WH>
 {
+    /// 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 +421,19 @@ impl<K: Hash + Eq, V> WTinyLFUCache<K, V, DefaultKeyHasher<K>> {
         )
         .finalize()
     }
+}
+
+impl<K: Hash + Eq, V, KH: KeyHasher<K>, FH: BuildHasher, RH: BuildHasher, WH: BuildHasher>
+    WTinyLFUCache<K, V, KH, FH, RH, WH>
+{
+    /// Creates a WTinyLFUCache according to [`WTinyLFUCacheBuilder`]
+    ///
+    /// [`WTinyLFUCacheBuilder`]: struct.WTinyLFUCacheBuilder.html
+    pub fn from_builder(
+        builder: WTinyLFUCacheBuilder<K, KH, FH, RH, WH>,
+    ) -> Result<Self, WTinyLFUError> {
+        builder.finalize()
+    }
 
     /// Returns a [`WTinyLFUCacheBuilder`] with default configurations.
     ///
@@ -429,19 +463,6 @@ impl<K: Hash + Eq, V> WTinyLFUCache<K, V, DefaultKeyHasher<K>> {
     }
 }
 
-impl<K: Hash + Eq, V, KH: KeyHasher<K>, FH: BuildHasher, RH: BuildHasher, WH: BuildHasher>
-    WTinyLFUCache<K, V, KH, FH, RH, WH>
-{
-    /// Creates a WTinyLFUCache according to [`WTinyLFUCacheBuilder`]
-    ///
-    /// [`WTinyLFUCacheBuilder`]: struct.WTinyLFUCacheBuilder.html
-    pub fn from_builder(
-        builder: WTinyLFUCacheBuilder<K, KH, FH, RH, WH>,
-    ) -> Result<Self, WTinyLFUError> {
-        builder.finalize()
-    }
-}
-
 impl<K: Hash + Eq, V, KH: KeyHasher<K>, FH: BuildHasher, RH: BuildHasher, WH: BuildHasher>
     Cache<K, V> for WTinyLFUCache<K, V, KH, FH, RH, WH>
 {
@@ -654,8 +675,107 @@ impl<K: Hash + Eq, V, KH: KeyHasher<K>, FH: BuildHasher, RH: BuildHasher, WH: Bu
 
 #[cfg(test)]
 mod test {
+    use core::hash::BuildHasher;
+
     use crate::lfu::WTinyLFUCache;
-    use crate::{Cache, PutResult};
+    use crate::lfu::tinylfu::test::{PassthroughU64KeyHasher, PassthroughU64Hasher};
+    use crate::{Cache, PutResult, WTinyLFUCacheBuilder};
+
+    #[derive(Default)]
+    struct PassthroughBuilderU64Hasher {}
+
+    impl BuildHasher for PassthroughBuilderU64Hasher {
+        type Hasher = PassthroughU64Hasher;
+
+        fn build_hasher(&self) -> Self::Hasher {
+            todo!()
+        }
+    }
+    #[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)]