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}