pbf/
pbf.rs

1// Copyright (c) 2023, Ruan Kunliang.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5use crate::hash::hash128;
6use std::cmp::max;
7
8const MAX_PAGE_NUM: usize = 1 << 18;
9
10pub trait BloomFilter {
11    fn get_way(&self) -> u8;
12    fn get_page_level(&self) -> u8;
13    fn get_page_num(&self) -> u32;
14    fn get_unique_cnt(&self) -> usize;
15    fn get_data(&self) -> &[u8];
16    fn capacity(&self) -> usize;
17    fn virtual_capacity(&self, fpr: f32) -> usize;
18    fn clear(&mut self);
19    fn valid(&self) -> bool;
20    fn set(&mut self, key: &[u8]) -> bool;
21    fn test(&self, key: &[u8]) -> bool;
22}
23
24pub const fn best_way_const(fpr: f32) -> u8 {
25    let mut rate = fpr;
26    if rate < 0.0005 {
27        rate = 0.0005;
28    } else if rate > 0.1 {
29        rate = 0.1;
30    }
31    // Approximate ceil(log2(2 / fpr)) with integer bit checks so this stays
32    // const-evaluable and matches the C++ selector.
33    let mut n = (2.0 / rate) as u32;
34    n += 1;
35
36    let mut i = 1_u32;
37    while i < 32 {
38        if (n & (0xfffffffe_u32 << i)) == 0 {
39            let way = i as u8;
40            if way < 4 {
41                return 4;
42            }
43            if way > 8 {
44                return 8;
45            }
46            return way;
47        }
48        i += 1;
49    }
50    0
51}
52
53fn estimate_params(item: usize, fpr: f32) -> (u8, u8, u32) {
54    let mut rate = fpr;
55    if rate < 0.0005 {
56        rate = 0.0005;
57    } else if rate > 0.1 {
58        rate = 0.1;
59    }
60    let way = best_way_const(rate);
61    let w = -f32::log2(rate);
62    let ln2 = f32::ln(2.0 as f32);
63    let mut bpi = (w / (ln2 * 8.0)) as f64;
64    if w > 9.0 {
65        let x = (w - 7.0) as f64;
66        bpi *= 1.0 + 0.0025*x*x;
67    } else if w > 3.0 {
68        bpi *= 1.01;
69    }
70
71    let n = (bpi * max(item, 1) as f64) as usize;
72    let mut page_level = 0_u8;
73    for i in 6..12 {
74        if n < (1_usize << (i + 4)) {
75            page_level = i;
76            if page_level < (8 - 8/way) {
77                page_level += 1;
78            }
79            break;
80        }
81    }
82    if page_level == 0 {
83        page_level = 12
84    }
85
86    let mut page_num = (n + (1 << page_level) - 1) >> page_level;
87    if page_num == 0 {
88        page_num = 1;
89    }
90    if page_num >= MAX_PAGE_NUM {
91        panic!("too many items");
92    }
93    (way, page_level, page_num as u32)
94}
95
96pub fn new_bloom_filter(item: usize, fpr: f32) -> Box<dyn BloomFilter> {
97    let (way, page_level, page_num) = estimate_params(item, fpr);
98    match way {
99        8 => Box::new(PageBloomFilter::<8>::new(page_level, page_num)),
100        7 => Box::new(PageBloomFilter::<7>::new(page_level, page_num)),
101        6 => Box::new(PageBloomFilter::<6>::new(page_level, page_num)),
102        5 => Box::new(PageBloomFilter::<5>::new(page_level, page_num)),
103        4 => Box::new(PageBloomFilter::<4>::new(page_level, page_num)),
104        _ => panic!("no way"),
105    }
106}
107
108pub fn new_pbf(way: u8, page_level: u8, page_num: u32) -> Box<dyn BloomFilter> {
109    match way {
110        8 => Box::new(PageBloomFilter::<8>::new(page_level, page_num)),
111        7 => Box::new(PageBloomFilter::<7>::new(page_level, page_num)),
112        6 => Box::new(PageBloomFilter::<6>::new(page_level, page_num)),
113        5 => Box::new(PageBloomFilter::<5>::new(page_level, page_num)),
114        4 => Box::new(PageBloomFilter::<4>::new(page_level, page_num)),
115        _ => panic!("no way"),
116    }
117}
118
119pub fn restore_pbf(way: u8, page_level: u8, data: &[u8], unique_cnt: usize) -> Box<dyn BloomFilter> {
120    match way {
121        8 => Box::new(PageBloomFilter::<8>::recover(page_level, data, unique_cnt)),
122        7 => Box::new(PageBloomFilter::<7>::recover(page_level, data, unique_cnt)),
123        6 => Box::new(PageBloomFilter::<6>::recover(page_level, data, unique_cnt)),
124        5 => Box::new(PageBloomFilter::<5>::recover(page_level, data, unique_cnt)),
125        4 => Box::new(PageBloomFilter::<4>::recover(page_level, data, unique_cnt)),
126        _ => panic!("no way"),
127    }
128}
129
130pub struct PageBloomFilter<const W : u8> {
131    page_level: u8,
132    page_num: u32,
133    unique_cnt: usize,
134    data: Vec<u8>,
135}
136
137impl<const W : u8> PageBloomFilter<W> {
138    pub fn from_estimate(item: usize, fpr: f32) -> Self {
139        let (way, page_level, page_num) = estimate_params(item, fpr);
140        assert_eq!(W, way, "generic W does not match estimated way");
141        Self::new(page_level, page_num)
142    }
143
144    pub fn new(page_level: u8, page_num: u32) -> Self {
145        if W < 4 || W > 8 {
146            panic!("way should be 4-8");
147        }
148        if page_level < (8-8/W) || page_level > 13 {
149            panic!("pageLevel should be 7-13");
150        }
151        if page_num == 0 || page_num as usize >= MAX_PAGE_NUM {
152            panic!("pageNum should be in [1, 1 << 18)");
153        }
154        return Self {
155            page_level: page_level,
156            page_num: page_num,
157            unique_cnt: 0,
158            data: vec![0_u8; (page_num as usize) << page_level],
159        };
160    }
161
162    pub fn recover(page_level: u8, data: &[u8], unique_cnt: usize) -> Self {
163        if W < 4 || W > 8 {
164            panic!("way should be 4-8");
165        }
166        if page_level < (8-8/W) || page_level > 13 {
167            panic!("pageLevel should be 7-13");
168        }
169        let page_size = 1_usize << page_level;
170        if data.len() == 0 || data.len()%page_size != 0 {
171            panic!("illegal data size");
172        }
173        let page_num = data.len() / page_size;
174        if page_num >= MAX_PAGE_NUM {
175            panic!("too big data");
176        }
177        return Self {
178            page_level: page_level,
179            page_num: page_num as u32,
180            unique_cnt: unique_cnt,
181            data: data.to_vec(),
182        };
183    }
184}
185
186#[inline(always)]
187fn rot(x: u32, k: u8) -> u32 {
188    return (x << k) | (x >> (32 - k));
189}
190
191#[inline(always)]
192fn page_hash(key: &[u8]) -> (u32, [u32; 4]) {
193    let code = hash128(key);
194    let w = [
195        code[0] as u32, (code[0] >> 32) as u32,
196        code[1] as u32, (code[1] >> 32) as u32,
197    ];
198    return (rot(w[0], 8) ^ rot(w[1], 6) ^ rot(w[2], 4) ^ rot(w[3], 2), w)
199}
200
201#[inline(always)]
202fn page_mask(page_level: u8) -> u16 {
203    ((1_u32 << (page_level + 3)) - 1) as u16
204}
205
206#[inline(always)]
207fn set_slot(data: &mut [u8], off: usize, mask: u16, word: u32, shift: u32, hit: &mut u8) {
208    let idx = ((word >> shift) as u16) & mask;
209    let byte = off + (idx >> 3) as usize;
210    let bit = 1_u8 << (idx & 7);
211    *hit &= data[byte] >> (idx & 7);
212    data[byte] |= bit;
213}
214
215#[inline(always)]
216fn test_slot(data: &[u8], off: usize, mask: u16, word: u32, shift: u32) -> bool {
217    let idx = ((word >> shift) as u16) & mask;
218    let bit = 1_u8 << (idx & 7);
219    (data[off + (idx >> 3) as usize] & bit) != 0
220}
221
222macro_rules! impl_bloom_filter {
223    ($w:literal, $(($word:literal, $shift:literal)),+ $(,)?) => {
224        impl BloomFilter for PageBloomFilter<$w> {
225            fn get_way(&self) -> u8 {
226                $w
227            }
228
229            fn get_page_level(&self) -> u8 {
230                self.page_level
231            }
232
233            fn get_page_num(&self) -> u32 {
234                self.page_num
235            }
236
237            fn get_unique_cnt(&self) -> usize {
238                self.unique_cnt
239            }
240
241            fn get_data(&self) -> &[u8] {
242                &self.data
243            }
244
245            fn capacity(&self) -> usize {
246                self.data.len() * 8 / $w as usize
247            }
248
249            fn virtual_capacity(&self, fpr: f32) -> usize {
250                let t = f64::ln_1p(-f64::powf(fpr as f64, 1.0 / $w as f64)) /
251                    f64::ln_1p(-1.0 / (self.data.len() * 8) as f64);
252                t as usize / $w as usize
253            }
254
255            fn clear(&mut self) {
256                self.unique_cnt = 0;
257                self.data.fill(0_u8);
258            }
259
260            fn valid(&self) -> bool {
261                !self.data.is_empty()
262            }
263
264            fn set(&mut self, key: &[u8]) -> bool {
265                let (code, words) = page_hash(key);
266                let off = ((code % self.page_num) as usize) << self.page_level;
267                let mask = page_mask(self.page_level);
268                let mut hit = 1_u8;
269                $(set_slot(&mut self.data, off, mask, words[$word], $shift, &mut hit);)+
270                if hit != 0 {
271                    return false;
272                }
273                self.unique_cnt += 1;
274                true
275            }
276
277            fn test(&self, key: &[u8]) -> bool {
278                let (code, words) = page_hash(key);
279                let off = ((code % self.page_num) as usize) << self.page_level;
280                let mask = page_mask(self.page_level);
281                $(if !test_slot(&self.data, off, mask, words[$word], $shift) {
282                    return false;
283                })+
284                true
285            }
286        }
287    };
288}
289
290impl_bloom_filter!(4, (0, 0), (0, 16), (1, 0), (1, 16));
291impl_bloom_filter!(5, (0, 0), (0, 16), (1, 0), (1, 16), (2, 0));
292impl_bloom_filter!(6, (0, 0), (0, 16), (1, 0), (1, 16), (2, 0), (2, 16));
293impl_bloom_filter!(7, (0, 0), (0, 16), (1, 0), (1, 16), (2, 0), (2, 16), (3, 0));
294impl_bloom_filter!(8, (0, 0), (0, 16), (1, 0), (1, 16), (2, 0), (2, 16), (3, 0), (3, 16));
pbf/pbf.rs

pbf/
pbf.rs
