Skip to content

Commit

Permalink
Make WTinyLFU / TinyLFU clonable
Browse files Browse the repository at this point in the history
  • Loading branch information
vlovich committed Oct 30, 2023
1 parent 7f23283 commit 2b44adb
Show file tree
Hide file tree
Showing 8 changed files with 135 additions and 0 deletions.
33 changes: 33 additions & 0 deletions src/lfu/tinylfu.rs
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,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.clone(),
w: self.w.clone(),
kh: self.kh.clone(),
marker: self.marker.clone(),
}
}
}

impl<K: Hash + Eq, KH: KeyHasher<K>> TinyLFU<K, KH> {
/// Returns a TinyLFU according to the [`TinyLFUBuilder`]
///
Expand Down Expand Up @@ -420,6 +433,26 @@ 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() {
Expand Down
1 change: 1 addition & 0 deletions src/lfu/tinylfu/bloom.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
Expand Down
1 change: 1 addition & 0 deletions src/lfu/tinylfu/sketch/count_min_row.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
1 change: 1 addition & 0 deletions src/lfu/tinylfu/sketch/count_min_sketch_std.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Expand Down
36 changes: 36 additions & 0 deletions src/lfu/wtinylfu.rs
Original file line number Diff line number Diff line change
Expand Up @@ -652,6 +652,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 crate::lfu::WTinyLFUCache;
Expand Down Expand Up @@ -741,4 +759,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(3, 3), PutResult::Evicted { key: 1, value: 2 });
}
}
1 change: 1 addition & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
Expand Down
49 changes: 49 additions & 0 deletions src/lru/raw.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
///
Expand Down Expand Up @@ -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));
}
}
13 changes: 13 additions & 0 deletions src/lru/segmented.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand Down

0 comments on commit 2b44adb

Please sign in to comment.