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