Skip to main content

protocache_core/
perfect_hash.rs

1//! Perfect-hash surface matching `perfect_hash.h`.
2
3pub use crate::serialize::{build_perfect_hash_index, build_perfect_hash_index_with_positions};
4
5use crate::hash::hash128;
6use crate::utils::{CorruptionKind, ReadError};
7
8#[derive(Clone, Copy, Debug)]
9pub struct PerfectHashView<'a> {
10    data: &'a [u8],
11    bitmap: &'a [u8],
12    table: &'a [u8],
13    section: usize,
14    section_magic: u64,
15    len: usize,
16    table_width: usize,
17}
18
19impl<'a> PerfectHashView<'a> {
20    #[inline(always)]
21    pub fn new(data: &'a [u8]) -> Result<Self, ReadError> {
22        let header = read_u32_le(data).ok_or(ReadError::new(CorruptionKind::Truncated))?;
23        let len = (header & 0x0fff_ffff) as usize;
24        if len <= 1 {
25            let data = data
26                .get(..4)
27                .ok_or(ReadError::new(CorruptionKind::Truncated))?;
28            return Ok(Self {
29                data,
30                bitmap: &data[..0],
31                table: &data[..0],
32                section: 0,
33                section_magic: 0,
34                len,
35                table_width: 0,
36            });
37        }
38
39        let (section, bitmap_size, table_width, bytes) = perfect_hash_layout(len);
40        let data = data
41            .get(..bytes)
42            .ok_or(ReadError::new(CorruptionKind::Truncated))?;
43        let bitmap = data
44            .get(8..8 + bitmap_size)
45            .ok_or(ReadError::new(CorruptionKind::Truncated))?;
46        let table = data
47            .get(8 + bitmap_size..)
48            .ok_or(ReadError::new(CorruptionKind::Truncated))?;
49
50        Ok(Self {
51            data,
52            bitmap,
53            table,
54            section,
55            section_magic: fast_mod_magic(section as u32),
56            len,
57            table_width,
58        })
59    }
60
61    #[inline(always)]
62    pub fn data_size(&self) -> usize {
63        self.data.len()
64    }
65
66    #[inline(always)]
67    pub fn len(&self) -> usize {
68        self.len
69    }
70
71    #[inline(always)]
72    pub fn is_empty(&self) -> bool {
73        self.len == 0
74    }
75
76    #[inline(always)]
77    pub fn locate(&self, key: &[u8]) -> Option<usize> {
78        if self.len == 0 {
79            return None;
80        }
81        if self.len == 1 {
82            return Some(0);
83        }
84
85        let seed = read_u32_le(self.data.get(4..8)?)? as u64;
86        let code = hash128(key, seed);
87        let slots = [
88            fast_mod_u32(code[0], self.section as u32, self.section_magic) as usize,
89            fast_mod_u32(code[1], self.section as u32, self.section_magic) as usize + self.section,
90            fast_mod_u32(code[2], self.section as u32, self.section_magic) as usize
91                + self.section * 2,
92        ];
93        self.locate_slots(slots)
94    }
95
96    #[inline(always)]
97    fn locate_slots(&self, slots: [usize; 3]) -> Option<usize> {
98        let m = bit2(self.bitmap, slots[0])?
99            + bit2(self.bitmap, slots[1])?
100            + bit2(self.bitmap, slots[2])?;
101        let slot = slots[(m % 3) as usize];
102        let a = slot >> 5;
103        let b = slot & 31;
104
105        let off = read_offset(self.table, self.table_width, a)?;
106
107        let block = read_u64_le(self.bitmap.get(a * 8..a * 8 + 8)?)? | (u64::MAX << (b << 1));
108        Some(off + count_valid_slot(block))
109    }
110}
111
112#[inline(always)]
113fn perfect_hash_layout(len: usize) -> (usize, usize, usize, usize) {
114    let section = ((len * 105).saturating_add(255) / 256).max(10);
115    let bitmap_size = ((section * 3 + 31) & !31) / 4;
116    let table_width = if len > u16::MAX as usize {
117        4
118    } else if len > u8::MAX as usize {
119        2
120    } else if len > 24 {
121        1
122    } else {
123        0
124    };
125    let bytes = 8 + bitmap_size + bitmap_size / 8 * table_width;
126    (section, bitmap_size, table_width, bytes)
127}
128
129#[inline(always)]
130fn read_offset(table: &[u8], table_width: usize, index: usize) -> Option<usize> {
131    match table_width {
132        4 => {
133            let start = index * 4;
134            Some(read_u32_le(table.get(start..start + 4)?)? as usize)
135        }
136        2 => {
137            let start = index * 2;
138            Some(u16::from_le_bytes(table.get(start..start + 2)?.try_into().ok()?) as usize)
139        }
140        1 => Some(*table.get(index)? as usize),
141        0 => Some(0),
142        _ => None,
143    }
144}
145
146#[inline(always)]
147fn read_u32_le(bytes: &[u8]) -> Option<u32> {
148    Some(u32::from_le_bytes(bytes.get(..4)?.try_into().ok()?))
149}
150
151#[inline(always)]
152fn read_u64_le(bytes: &[u8]) -> Option<u64> {
153    Some(u64::from_le_bytes(bytes.get(..8)?.try_into().ok()?))
154}
155
156#[inline(always)]
157fn bit2(vec: &[u8], pos: usize) -> Option<u32> {
158    Some(((vec.get(pos >> 2)? >> ((pos & 3) << 1)) & 3) as u32)
159}
160
161#[inline(always)]
162fn fast_mod_magic(divisor: u32) -> u64 {
163    u64::MAX / divisor as u64 + 1
164}
165
166#[inline(always)]
167fn fast_mod_u32(value: u32, divisor: u32, magic: u64) -> u32 {
168    let low = magic.wrapping_mul(value as u64);
169    (((low as u128) * divisor as u128) >> 64) as u32
170}
171
172#[inline(always)]
173fn count_valid_slot(v: u64) -> usize {
174    let invalid = ((v & 0x5555_5555_5555_5555) & (v >> 1)).count_ones() as usize;
175    32 - invalid
176}
177
178#[cfg(test)]
179mod tests {
180    use std::collections::HashSet;
181
182    use super::{PerfectHashView, build_perfect_hash_index, fast_mod_magic, fast_mod_u32};
183
184    fn run_case(size: usize) {
185        let keys = (0..size).map(|i| i.to_string()).collect::<Vec<_>>();
186        let index = build_perfect_hash_index(&keys).expect("perfect hash index should build");
187        let view = PerfectHashView::new(&index).expect("perfect hash view should decode");
188
189        assert_eq!(view.len(), keys.len());
190        assert_eq!(view.data_size(), index.len());
191
192        let mut seen = HashSet::new();
193        for key in &keys {
194            let pos = view
195                .locate(key.as_bytes())
196                .expect("every inserted key should resolve");
197            assert!(pos < keys.len());
198            assert!(seen.insert(pos), "duplicate slot {pos} for key {key}");
199        }
200    }
201
202    #[test]
203    fn tiny_cases_match_cpp_coverage() {
204        for size in [0usize, 1, 2, 24] {
205            run_case(size);
206        }
207    }
208
209    #[test]
210    fn small_cases_match_cpp_coverage() {
211        for size in [200usize, 255, 1000] {
212            run_case(size);
213        }
214    }
215
216    #[test]
217    fn big_cases_match_cpp_coverage() {
218        for size in [65_535usize, 100_000] {
219            run_case(size);
220        }
221    }
222
223    #[test]
224    fn rejects_truncated_indices() {
225        assert!(PerfectHashView::new(&[]).is_err());
226        assert!(PerfectHashView::new(&[0, 0, 0]).is_err());
227
228        let index = build_perfect_hash_index(&["a", "b", "c"]).unwrap();
229        for len in 0..index.len() {
230            assert!(PerfectHashView::new(&index[..len]).is_err());
231        }
232    }
233
234    #[test]
235    fn fast_mod_matches_remainder() {
236        let divisors = [10u32, 11, 17, 31, 45, 255, 256, 1000, 65535, 100_000];
237        let values = [
238            0u32,
239            1,
240            2,
241            3,
242            7,
243            31,
244            32,
245            255,
246            256,
247            1024,
248            65_535,
249            1_000_000,
250            u32::MAX - 1,
251            u32::MAX,
252        ];
253
254        for divisor in divisors {
255            let magic = fast_mod_magic(divisor);
256            for value in values {
257                assert_eq!(fast_mod_u32(value, divisor, magic), value % divisor);
258            }
259        }
260    }
261}