Skip to main content

harfrust/hb/
cache.rs

1use core::sync::atomic::{AtomicU16, AtomicU32, AtomicU8, Ordering};
2
3/// Trait for atomics used in cache storage
4pub trait AtomicStorage: Sized {
5    const BITS: usize;
6    fn get(&self) -> u32;
7    fn set(&self, val: u32);
8    fn default() -> Self;
9}
10
11impl AtomicStorage for AtomicU8 {
12    const BITS: usize = 8;
13
14    fn get(&self) -> u32 {
15        self.load(Ordering::Relaxed) as u32
16    }
17
18    fn set(&self, val: u32) {
19        self.store(val as u8, Ordering::Relaxed);
20    }
21
22    fn default() -> Self {
23        Self::new(u8::MAX)
24    }
25}
26
27impl AtomicStorage for AtomicU16 {
28    const BITS: usize = 16;
29
30    fn get(&self) -> u32 {
31        self.load(Ordering::Relaxed) as u32
32    }
33
34    fn set(&self, val: u32) {
35        self.store(val as u16, Ordering::Relaxed);
36    }
37
38    fn default() -> Self {
39        Self::new(u16::MAX)
40    }
41}
42
43impl AtomicStorage for AtomicU32 {
44    const BITS: usize = 32;
45
46    fn get(&self) -> u32 {
47        self.load(Ordering::Relaxed)
48    }
49
50    fn set(&self, val: u32) {
51        self.store(val, Ordering::Relaxed);
52    }
53
54    fn default() -> Self {
55        Self::new(u32::MAX)
56    }
57}
58
59/// Selects correct type from STORAGE_BITS
60pub trait SelectAtomic<const BITS: usize> {
61    type Type: AtomicStorage;
62}
63impl SelectAtomic<8> for () {
64    type Type = AtomicU8;
65}
66impl SelectAtomic<16> for () {
67    type Type = AtomicU16;
68}
69impl SelectAtomic<32> for () {
70    type Type = AtomicU32;
71}
72
73/// Public wrapper
74pub type hb_cache_t<
75    const KEY_BITS: usize,
76    const VALUE_BITS: usize,
77    const CACHE_SIZE: usize,
78    const STORAGE_BITS: usize,
79> = hb_cache_core_t<KEY_BITS, VALUE_BITS, CACHE_SIZE, <() as SelectAtomic<STORAGE_BITS>>::Type>;
80
81/// Core cache
82#[derive(Debug)]
83pub struct hb_cache_core_t<
84    const KEY_BITS: usize,
85    const VALUE_BITS: usize,
86    const CACHE_SIZE: usize,
87    T: AtomicStorage,
88> {
89    values: [T; CACHE_SIZE],
90}
91
92impl<const KEY_BITS: usize, const VALUE_BITS: usize, const CACHE_SIZE: usize, T: AtomicStorage>
93    hb_cache_core_t<KEY_BITS, VALUE_BITS, CACHE_SIZE, T>
94{
95    pub const MAX_VALUE: u32 = (1 << VALUE_BITS) - 1;
96    const CACHE_BITS: usize = CACHE_SIZE.ilog2() as usize;
97
98    pub fn new() -> Self {
99        debug_assert!(
100            CACHE_SIZE.is_power_of_two(),
101            "CACHE_SIZE must be a power of two"
102        );
103
104        debug_assert!(
105            KEY_BITS >= Self::CACHE_BITS,
106            "KEY_BITS must be >= log2(CACHE_SIZE)"
107        );
108        debug_assert!(
109            KEY_BITS + VALUE_BITS <= Self::CACHE_BITS + T::BITS,
110            "KEY_BITS + VALUE_BITS must fit in CACHE_BITS + T::BITS"
111        );
112
113        Self {
114            values: core::array::from_fn(|_| T::default()),
115        }
116    }
117
118    #[inline]
119    pub fn get(&self, key: u32) -> Option<u32> {
120        let index = (key as usize) & (CACHE_SIZE - 1);
121        let stored = self.values[index].get();
122        let tag = stored >> VALUE_BITS;
123        let expected_tag = key >> Self::CACHE_BITS;
124
125        if stored == T::default().get() || tag != expected_tag {
126            return None;
127        }
128
129        Some(stored & ((1 << VALUE_BITS) - 1))
130    }
131
132    #[inline]
133    pub fn set(&self, key: u32, value: u32) {
134        if (key >> KEY_BITS) != 0 || (value >> VALUE_BITS) != 0 {
135            return;
136        }
137        self.set_unchecked(key, value);
138    }
139
140    #[inline]
141    fn set_unchecked(&self, key: u32, value: u32) {
142        let index = (key as usize) & (CACHE_SIZE - 1);
143        let packed = ((key >> Self::CACHE_BITS) << VALUE_BITS) | value;
144        self.values[index].set(packed);
145    }
146}