pgm_extra/util/
cache.rs

1use alloc::vec::Vec;
2use core::cell::Cell;
3
4const CACHE_SIZE: usize = 64;
5
6#[derive(Clone, Copy)]
7struct CacheEntry<K: Copy> {
8    key: K,
9    pos: usize,
10    valid: bool,
11}
12
13impl<K: Copy + Default> Default for CacheEntry<K> {
14    fn default() -> Self {
15        Self {
16            key: K::default(),
17            pos: 0,
18            valid: false,
19        }
20    }
21}
22
23pub trait FastHash: Copy + Eq {
24    fn fast_hash(&self) -> usize;
25}
26
27macro_rules! impl_fast_hash_int {
28    ($($t:ty),*) => {
29        $(
30            impl FastHash for $t {
31                #[inline(always)]
32                fn fast_hash(&self) -> usize {
33                    let x = *self as u64;
34                    let x = x.wrapping_mul(0x9E3779B97F4A7C15);
35                    (x >> 58) as usize
36                }
37            }
38        )*
39    };
40}
41
42impl_fast_hash_int!(u8, u16, u32, u64, usize, i8, i16, i32, i64, isize);
43
44impl FastHash for u128 {
45    #[inline(always)]
46    fn fast_hash(&self) -> usize {
47        let x = (*self as u64).wrapping_mul(0x9E3779B97F4A7C15);
48        (x >> 58) as usize
49    }
50}
51
52impl FastHash for i128 {
53    #[inline(always)]
54    fn fast_hash(&self) -> usize {
55        let x = (*self as u64).wrapping_mul(0x9E3779B97F4A7C15);
56        (x >> 58) as usize
57    }
58}
59
60pub struct HotCache<K: Copy + Default + FastHash> {
61    entries: Vec<Cell<CacheEntry<K>>>,
62}
63
64impl<K: Copy + Default + FastHash> core::fmt::Debug for HotCache<K> {
65    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
66        f.debug_struct("HotCache")
67            .field("size", &self.entries.len())
68            .finish()
69    }
70}
71
72impl<K: Copy + Default + FastHash> HotCache<K> {
73    pub fn new() -> Self {
74        let mut entries = Vec::with_capacity(CACHE_SIZE);
75        for _ in 0..CACHE_SIZE {
76            entries.push(Cell::new(CacheEntry::default()));
77        }
78        Self { entries }
79    }
80
81    #[inline(always)]
82    fn hash_key(key: &K) -> usize {
83        key.fast_hash()
84    }
85
86    #[inline]
87    pub fn lookup(&self, key: &K) -> Option<usize> {
88        let idx = Self::hash_key(key);
89        if idx >= self.entries.len() {
90            return None;
91        }
92
93        let entry = self.entries[idx].get();
94
95        if entry.valid && entry.key == *key {
96            Some(entry.pos)
97        } else {
98            None
99        }
100    }
101
102    #[inline]
103    pub fn insert(&self, key: K, pos: usize) {
104        let idx = Self::hash_key(&key);
105        if idx >= self.entries.len() {
106            return;
107        }
108
109        self.entries[idx].set(CacheEntry {
110            key,
111            pos,
112            valid: true,
113        });
114    }
115
116    pub fn clear(&self) {
117        for cell in &self.entries {
118            let mut entry = cell.get();
119            entry.valid = false;
120            cell.set(entry);
121        }
122    }
123}
124
125impl<K: Copy + Default + FastHash> Default for HotCache<K> {
126    fn default() -> Self {
127        Self::new()
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134
135    #[test]
136    fn test_cache_basic() {
137        let cache: HotCache<u64> = HotCache::new();
138
139        assert_eq!(cache.lookup(&42), None);
140
141        cache.insert(42, 100);
142        assert_eq!(cache.lookup(&42), Some(100));
143
144        cache.insert(42, 200);
145        assert_eq!(cache.lookup(&42), Some(200));
146    }
147
148    #[test]
149    fn test_cache_miss() {
150        let cache: HotCache<u64> = HotCache::new();
151        cache.insert(42, 100);
152
153        assert_eq!(cache.lookup(&43), None);
154    }
155
156    #[test]
157    fn test_cache_collision() {
158        let cache: HotCache<u64> = HotCache::new();
159
160        cache.insert(0, 100);
161        cache.insert(64, 200);
162
163        let hit0 = cache.lookup(&0);
164        let hit64 = cache.lookup(&64);
165
166        assert!(hit0.is_some() || hit64.is_some());
167    }
168
169    #[test]
170    fn test_cache_clear() {
171        let cache: HotCache<u64> = HotCache::new();
172        cache.insert(42, 100);
173        cache.clear();
174        assert_eq!(cache.lookup(&42), None);
175    }
176}