quill_sql/utils/cache/
tiny_lfu.rs

1use std::collections::hash_map::DefaultHasher;
2use std::hash::{Hash, Hasher};
3
4/// TinyLFU-style admission filter: approximate frequency via 4-bit counters in a simple
5/// Count-Min Sketch. This is a minimal, lockless (external locking) and CPU-cheap version
6/// intended to bias admission decisions before the main replacer.
7#[derive(Debug)]
8pub struct TinyLFU {
9    width: usize,
10    depth: usize,
11    tables: Vec<Vec<u8>>, // 4-bit per counter packed into u8 (2 counters per byte)
12}
13
14impl TinyLFU {
15    pub fn new(width: usize, depth: usize) -> Self {
16        let width = width.next_power_of_two();
17        let depth = depth.max(1).min(4);
18        let tables = (0..depth).map(|_| vec![0u8; (width + 1) / 2]).collect();
19        Self {
20            width,
21            depth,
22            tables,
23        }
24    }
25
26    #[inline]
27    fn hash_i(&self, key: u64, i: usize) -> usize {
28        let mut h = DefaultHasher::new();
29        (key.wrapping_add((i as u64) << 32)).hash(&mut h);
30        (h.finish() as usize) & (self.width - 1)
31    }
32
33    #[inline]
34    fn load_counter(slot: &mut [u8], idx: usize) -> u8 {
35        let byte = &mut slot[idx / 2];
36        if idx % 2 == 0 {
37            *byte & 0x0F
38        } else {
39            (*byte >> 4) & 0x0F
40        }
41    }
42
43    #[inline]
44    fn store_counter(slot: &mut [u8], idx: usize, val: u8) {
45        let b = &mut slot[idx / 2];
46        if idx % 2 == 0 {
47            *b = (*b & 0xF0) | (val & 0x0F);
48        } else {
49            *b = (*b & 0x0F) | ((val & 0x0F) << 4);
50        }
51    }
52
53    /// Record an access for the 64-bit key.
54    pub fn admit_record(&mut self, key: u64) {
55        for i in 0..self.depth {
56            let idx = self.hash_i(key, i);
57            let slot = &mut self.tables[i];
58            let mut c = Self::load_counter(slot, idx);
59            if c < 15 {
60                c += 1;
61            }
62            Self::store_counter(slot, idx, c);
63        }
64    }
65
66    /// Estimate frequency for the key (min of counters).
67    pub fn estimate(&self, key: u64) -> u8 {
68        let mut minv = 15u8;
69        for i in 0..self.depth {
70            let idx = self.hash_i(key, i);
71            let slot = &self.tables[i];
72            let c = if idx / 2 < slot.len() {
73                if idx % 2 == 0 {
74                    slot[idx / 2] & 0x0F
75                } else {
76                    (slot[idx / 2] >> 4) & 0x0F
77                }
78            } else {
79                0
80            };
81            if c < minv {
82                minv = c;
83            }
84        }
85        minv
86    }
87
88    /// Periodic aging to prevent counter saturation. Halves all counters.
89    pub fn age(&mut self) {
90        for t in self.tables.iter_mut() {
91            for b in t.iter_mut() {
92                let lo = (*b & 0x0F) >> 1;
93                let hi = ((*b >> 4) & 0x0F) >> 1;
94                *b = (hi << 4) | lo;
95            }
96        }
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[test]
105    fn tiny_lfu_basic() {
106        let mut f = TinyLFU::new(1024, 4);
107        let k1 = 123u64;
108        let k2 = 456u64;
109        for _ in 0..8 {
110            f.admit_record(k1);
111        }
112        for _ in 0..2 {
113            f.admit_record(k2);
114        }
115        assert!(f.estimate(k1) >= f.estimate(k2));
116        f.age();
117        assert!(f.estimate(k1) >= f.estimate(k2));
118    }
119}