k_cache/
cache.rs

1use std::hash::Hash;
2use std::marker::PhantomData;
3use std::sync::atomic::AtomicBool;
4use std::sync::Arc;
5use std::{
6    borrow::Borrow,
7    collections::{HashMap, VecDeque},
8    hash::BuildHasher,
9};
10
11pub trait Weigher<K, V> {
12    fn weigh(_k: &K, _v: &V) -> usize {
13        1
14    }
15}
16
17#[derive(Debug, Clone)]
18pub struct One;
19impl<K, V> Weigher<K, V> for One {}
20
21struct SieveEntry<D> {
22    data: D,
23    visited: Arc<AtomicBool>,
24}
25
26pub struct Cache<K, V, S, W: Weigher<K, V> = One> {
27    map: HashMap<K, SieveEntry<V>, S>,
28    sieve_pool: VecDeque<SieveEntry<K>>,
29    sieve_hand: usize,
30    max_weight: usize,
31    weight: usize,
32    _phantom: PhantomData<W>,
33}
34
35impl<K, V, S, W> Cache<K, V, S, W>
36where
37    K: Eq + Hash + Clone,
38    S: BuildHasher,
39    W: Weigher<K, V>,
40{
41    pub fn new(hasher: S, max_weight: usize) -> Self {
42        Self {
43            map: HashMap::with_hasher(hasher),
44            sieve_pool: VecDeque::new(),
45            sieve_hand: 0,
46            max_weight,
47            weight: 0,
48            _phantom: PhantomData,
49        }
50    }
51
52    pub fn put(&mut self, key: K, value: V) {
53        let new_weight = self.make_room_for(&key, &value);
54        self.weight += new_weight;
55        let visited = Arc::new(AtomicBool::new(true));
56        if let Some(replaced) = self.map.insert(
57            key.clone(),
58            SieveEntry {
59                data: value,
60                visited: visited.clone(),
61            },
62        ) {
63            let replaced_weight = W::weigh(&key, &replaced.data);
64            match self.weight.checked_sub(replaced_weight) {
65                Some(new_weight) => self.weight = new_weight,
66                None => {
67                    log::error!("weight underflow");
68                    self.weight = 0;
69                }
70            }
71        }
72        self.sieve_pool.push_back(SieveEntry {
73            data: key.clone(),
74            visited,
75        });
76    }
77
78    pub fn get<Q>(&self, key: &Q) -> Option<&V>
79    where
80        K: Borrow<Q>,
81        Q: Hash + Eq + ?Sized,
82    {
83        match self.map.get(key) {
84            Some(entry) => {
85                entry
86                    .visited
87                    .store(true, std::sync::atomic::Ordering::Relaxed);
88                Some(&entry.data)
89            }
90            None => None,
91        }
92    }
93
94    fn make_room_for(&mut self, key: &K, value: &V) -> usize {
95        let entry_weight = W::weigh(key, value);
96        while self.max_weight < self.weight + entry_weight {
97            let sieve_entry = &mut self.sieve_pool[self.sieve_hand];
98            let visited = sieve_entry
99                .visited
100                .swap(false, std::sync::atomic::Ordering::Relaxed);
101            if visited {
102                self.sieve_hand = (self.sieve_hand + 1) % self.sieve_pool.len();
103            } else {
104                let sieve_key_entry = self
105                    .sieve_pool
106                    .swap_remove_back(self.sieve_hand)
107                    .expect("the index must be present");
108                match self.map.remove(&sieve_key_entry.data) {
109                    Some(removed) => {
110                        let removed_weight = W::weigh(&sieve_key_entry.data, &removed.data);
111                        match self.weight.checked_sub(removed_weight) {
112                            Some(new_weight) => self.weight = new_weight,
113                            None => {
114                                log::error!("weight underflow");
115                                self.weight = 0;
116                            }
117                        }
118                    }
119                    None => {
120                        log::debug!("garbage collecting sieve entry as {}", self.sieve_hand);
121                    }
122                }
123
124                if self.sieve_hand == self.sieve_pool.len() {
125                    self.sieve_hand = 0;
126                }
127            }
128        }
129        entry_weight
130    }
131}
k_cache/cache.rs

k_cache/
cache.rs
