From 07465913b1771d755a2c305c6f6d75f00b2e477c Mon Sep 17 00:00:00 2001
From: Vitali Lovich <vlovich+github@gmail.com>
Date: Tue, 31 Oct 2023 10:19:38 -0700
Subject: [PATCH] Make WTinyLFU / TinyLFU clonable (#19)

---
 src/lfu/tinylfu.rs                            | 33 +++++++++++++
 src/lfu/tinylfu/bloom.rs                      |  1 +
 src/lfu/tinylfu/sketch/count_min_row.rs       |  1 +
 .../tinylfu/sketch/count_min_sketch_std.rs    |  1 +
 src/lfu/wtinylfu.rs                           | 36 ++++++++++++++
 src/lib.rs                                    |  1 +
 src/lru/raw.rs                                | 49 +++++++++++++++++++
 src/lru/segmented.rs                          | 13 +++++
 8 files changed, 135 insertions(+)

diff --git a/src/lfu/tinylfu.rs b/src/lfu/tinylfu.rs
index e57208d..9d05128 100644
--- a/src/lfu/tinylfu.rs
+++ b/src/lfu/tinylfu.rs
@@ -150,6 +150,19 @@ impl<K: Hash + Eq> TinyLFU<K> {
     }
 }
 
+impl<K: Clone, KH: Clone> Clone for TinyLFU<K, KH> {
+    fn clone(&self) -> Self {
+        Self {
+            ctr: self.ctr.clone(),
+            doorkeeper: self.doorkeeper.clone(),
+            samples: self.samples,
+            w: self.w,
+            kh: self.kh.clone(),
+            marker: self.marker,
+        }
+    }
+}
+
 impl<K: Hash + Eq, KH: KeyHasher<K>> TinyLFU<K, KH> {
     /// Returns a TinyLFU according to the [`TinyLFUBuilder`]
     ///
@@ -488,6 +501,26 @@ pub(crate) mod test {
         assert_eq!(l.ctr.estimate(1), 1);
     }
 
+    #[test]
+    fn test_clone() {
+        let mut l: TinyLFU<u64> = TinyLFU::new(4, 4, 0.01).unwrap();
+        l.increment_hashed_key(1);
+        l.increment_hashed_key(1);
+        l.increment_hashed_key(1);
+
+        assert!(l.doorkeeper.contains(1));
+        assert_eq!(l.ctr.estimate(1), 2);
+
+        let cloned = l.clone();
+
+        l.increment_hashed_key(1);
+        assert!(!l.doorkeeper.contains(1));
+        assert_eq!(l.ctr.estimate(1), 1);
+
+        assert!(cloned.doorkeeper.contains(1));
+        assert_eq!(cloned.ctr.estimate(1), 2);
+    }
+
     // TODO: fix the bug caused by random
     // #[test]
     // fn test_estimate() {
diff --git a/src/lfu/tinylfu/bloom.rs b/src/lfu/tinylfu/bloom.rs
index 25a7319..e22ed4e 100644
--- a/src/lfu/tinylfu/bloom.rs
+++ b/src/lfu/tinylfu/bloom.rs
@@ -46,6 +46,7 @@ fn calc_size_by_wrong_positives(num_entries: usize, wrongs: f64) -> EntriesLocs
 }
 
 /// Bloom filter
+#[derive(Clone)]
 #[repr(C)]
 pub(crate) struct Bloom {
     bitset: Vec<u64>,
diff --git a/src/lfu/tinylfu/sketch/count_min_row.rs b/src/lfu/tinylfu/sketch/count_min_row.rs
index e2f7217..3899c7e 100644
--- a/src/lfu/tinylfu/sketch/count_min_row.rs
+++ b/src/lfu/tinylfu/sketch/count_min_row.rs
@@ -10,6 +10,7 @@ use alloc::vec::Vec;
 use core::fmt::{Debug, Formatter};
 use core::ops::{Index, IndexMut};
 
+#[derive(Clone)]
 pub(crate) struct CountMinRow(Vec<u8>);
 
 impl CountMinRow {
diff --git a/src/lfu/tinylfu/sketch/count_min_sketch_std.rs b/src/lfu/tinylfu/sketch/count_min_sketch_std.rs
index 14afeff..70e0abb 100644
--- a/src/lfu/tinylfu/sketch/count_min_sketch_std.rs
+++ b/src/lfu/tinylfu/sketch/count_min_sketch_std.rs
@@ -12,6 +12,7 @@ use rand::{rngs::StdRng, Rng, SeedableRng};
 
 /// `CountMinSketch` is a small conservative-update count-min sketch
 /// implementation with 4-bit counters
+#[derive(Clone)]
 pub(crate) struct CountMinSketch {
     rows: [CountMinRow; DEPTH],
     seeds: [u64; DEPTH],
diff --git a/src/lfu/wtinylfu.rs b/src/lfu/wtinylfu.rs
index b6f8ba5..aebb831 100644
--- a/src/lfu/wtinylfu.rs
+++ b/src/lfu/wtinylfu.rs
@@ -687,6 +687,24 @@ impl<K: Hash + Eq, V, KH: KeyHasher<K>, FH: BuildHasher, RH: BuildHasher, WH: Bu
     }
 }
 
+impl<
+        K: Hash + Eq + Clone,
+        V: Clone,
+        KH: KeyHasher<K> + Clone,
+        FH: BuildHasher + Clone,
+        RH: BuildHasher + Clone,
+        WH: BuildHasher + Clone,
+    > Clone for WTinyLFUCache<K, V, KH, FH, RH, WH>
+{
+    fn clone(&self) -> Self {
+        Self {
+            tinylfu: self.tinylfu.clone(),
+            lru: self.lru.clone(),
+            slru: self.slru.clone(),
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use core::hash::BuildHasher;
@@ -888,4 +906,22 @@ mod test {
         assert_eq!(cache.remove(&3), Some(33));
         assert_eq!(cache.remove(&2), Some(22));
     }
+
+    #[test]
+    #[cfg_attr(miri, ignore)]
+    fn test_wtinylfu_clone() {
+        let mut cache = WTinyLFUCache::with_sizes(1, 2, 2, 5).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));
+
+        assert_eq!(cache.put(4, 3), PutResult::Evicted { key: 1, value: 1 });
+    }
 }
diff --git a/src/lib.rs b/src/lib.rs
index fcb8a2c..af6ca60 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -134,6 +134,7 @@ cfg_std!(
 
 // Struct used to hold a reference to a key
 #[doc(hidden)]
+#[derive(Clone)]
 pub struct KeyRef<K> {
     k: *const K,
 }
diff --git a/src/lru/raw.rs b/src/lru/raw.rs
index e3474c2..42a713d 100644
--- a/src/lru/raw.rs
+++ b/src/lru/raw.rs
@@ -198,6 +198,29 @@ pub struct RawLRU<K, V, E = DefaultEvictCallback, S = DefaultHashBuilder> {
     tail: *mut EntryNode<K, V>,
 }
 
+impl<K: Hash + Eq + Clone, V: Clone, E: OnEvictCallback + Clone, S: BuildHasher + Clone> Clone
+    for RawLRU<K, V, E, S>
+{
+    fn clone(&self) -> Self {
+        let mut cloned = RawLRU::<K, V, E, S>::construct(
+            self.cap,
+            HashMap::with_capacity_and_hasher(self.map.capacity(), self.map.hasher().clone()),
+            self.on_evict.clone(),
+        );
+        for entry in self.map.values() {
+            let (k, v) = unsafe {
+                let entry = entry.as_ref();
+                (
+                    entry.key.assume_init_ref().clone(),
+                    entry.val.assume_init_ref().clone(),
+                )
+            };
+            cloned.put(k, v);
+        }
+        cloned
+    }
+}
+
 impl<K: Hash + Eq, V> RawLRU<K, V> {
     /// Creates a new LRU Cache that holds at most `cap` items.
     ///
@@ -3032,4 +3055,30 @@ mod tests {
         let cache = RawLRU::<u64, u64>::new(0);
         assert_eq!(cache.unwrap_err(), CacheError::InvalidSize(0))
     }
+
+    #[test]
+    fn test_clone() {
+        let mut cache = RawLRU::<u64, u64>::new(2).unwrap();
+        cache.put(5, 6);
+        assert_eq!(cache.len(), 1);
+        assert_eq!(cache.get(&5), Some(&6));
+
+        let mut clone = cache.clone();
+        assert_eq!(clone.len(), 1);
+        assert_eq!(clone.get(&5), Some(&6));
+
+        cache.put(6, 7);
+        assert_eq!(cache.len(), 2);
+        assert_eq!(clone.len(), 1);
+
+        clone.put(1, 2);
+        assert_eq!(cache.len(), 2);
+        assert_eq!(clone.len(), 2);
+
+        std::mem::drop(cache);
+
+        clone.put(2, 3);
+        assert_eq!(clone.len(), 2);
+        assert_eq!(clone.peek(&2), Some(&3));
+    }
 }
diff --git a/src/lru/segmented.rs b/src/lru/segmented.rs
index 3bf7284..f9e1a37 100644
--- a/src/lru/segmented.rs
+++ b/src/lru/segmented.rs
@@ -157,6 +157,19 @@ pub struct SegmentedCache<K, V, FH = DefaultHashBuilder, RH = DefaultHashBuilder
     protected: RawLRU<K, V, DefaultEvictCallback, FH>,
 }
 
+impl<K: Hash + Eq + Clone, V: Clone, FH: BuildHasher + Clone, RH: BuildHasher + Clone> Clone
+    for SegmentedCache<K, V, FH, RH>
+{
+    fn clone(&self) -> Self {
+        Self {
+            probationary_size: self.probationary_size,
+            probationary: self.probationary.clone(),
+            protected_size: self.protected_size,
+            protected: self.protected.clone(),
+        }
+    }
+}
+
 impl<K: Hash + Eq, V> SegmentedCache<K, V> {
     /// Create a `SegmentedCache` with size and default configurations.
     pub fn new(probationary_size: usize, protected_size: usize) -> Result<Self, CacheError> {