k_cache/
segmented.rs

1use std::{borrow::Borrow, hash::BuildHasher};
2
3use crate::{Cache, One, Weigher};
4
5pub struct SegmentedCache<K, V, S: BuildHasher = std::hash::RandomState, W: Weigher<K, V> = One> {
6    segments: Vec<k_lock::Mutex<Cache<K, V, S, W>>>,
7    hasher: S,
8}
9
10impl<K, V, S, W> SegmentedCache<K, V, S, W>
11where
12    K: Eq + std::hash::Hash + Clone,
13    V: Clone,
14    S: BuildHasher + Default,
15    W: Weigher<K, V> + Clone,
16{
17    pub fn new(segments: usize, max_weight: usize) -> Self {
18        let weight_per_segment = max_weight / segments;
19        let segments = (0..segments)
20            .map(|_| k_lock::Mutex::new(Cache::new(S::default(), weight_per_segment)))
21            .collect();
22        Self {
23            segments,
24            hasher: S::default(),
25        }
26    }
27
28    pub fn put(&self, key: K, value: V) {
29        let slot = self.hasher.hash_one(&key) as usize % self.segments.len();
30        self.segments[slot]
31            .lock()
32            .expect("mutex must not be poisoned")
33            .put(key, value)
34    }
35
36    pub fn get<Q>(&self, key: &Q) -> Option<V>
37    where
38        K: Borrow<Q>,
39        Q: std::hash::Hash + Eq + ?Sized,
40    {
41        let slot = self.hasher.hash_one(key) as usize % self.segments.len();
42        self.segments[slot]
43            .lock()
44            .expect("mutex must not be poisoned")
45            .get(key)
46            .cloned()
47    }
48}