1#![allow(
7 clippy::cast_lossless,
8 clippy::cast_possible_truncation,
9 clippy::cast_possible_wrap,
10 clippy::cast_precision_loss,
11 clippy::cast_sign_loss,
12 clippy::doc_markdown,
13 clippy::items_after_statements,
14 clippy::similar_names,
15 clippy::unreadable_literal
16)]
17
18use alloc::format;
51use alloc::string::String;
52use alloc::vec;
53use alloc::vec::Vec;
54use core::fmt;
55
56use spg_crypto::crc32::crc32;
57
58const BLOOM_MAGIC: u32 = 0xB100_F11E;
62
63const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
65const FNV_PRIME: u64 = 0x0000_0001_0000_01b3;
66
67const NUM_HASHES_MAX: u32 = 32;
71
72#[derive(Debug, PartialEq, Eq)]
76pub enum BloomError {
77 TooShort { got: usize, need: usize },
79 BadMagic { got: u32 },
81 BadShape(String),
84 BadCrc { expected: u32, got: u32 },
86 BadNumHashes { got: u32 },
88}
89
90impl fmt::Display for BloomError {
91 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
92 match self {
93 Self::TooShort { got, need } => {
94 write!(f, "bloom: too short, got {got} bytes, need at least {need}")
95 }
96 Self::BadMagic { got } => {
97 write!(
98 f,
99 "bloom: bad magic 0x{got:08x}, expected 0x{BLOOM_MAGIC:08x}"
100 )
101 }
102 Self::BadShape(s) => write!(f, "bloom: bad shape: {s}"),
103 Self::BadCrc { expected, got } => write!(
104 f,
105 "bloom: crc mismatch, expected 0x{expected:08x}, got 0x{got:08x}"
106 ),
107 Self::BadNumHashes { got } => write!(
108 f,
109 "bloom: bad num_hashes {got}, must be 1..={NUM_HASHES_MAX}"
110 ),
111 }
112 }
113}
114
115#[derive(Debug, Clone)]
120pub struct BloomFilter {
121 bits: Vec<u64>,
122 num_bits: u64,
126 num_hashes: u32,
127}
128
129impl BloomFilter {
130 #[must_use]
148 pub fn with_target_fp_rate(num_items: usize, fp_rate: f64) -> Self {
149 assert!(num_items > 0, "BloomFilter: num_items must be > 0");
150 assert!(
151 fp_rate > 0.0 && fp_rate < 1.0,
152 "BloomFilter: fp_rate must be in (0, 1), got {fp_rate}"
153 );
154 let n = num_items as f64;
160 let ln_2 = libm_ln(2.0);
161 let m_raw = -(n * libm_ln(fp_rate)) / (ln_2 * ln_2);
162 let m_ceil_bits = f64_ceil_to_u64(m_raw).max(64);
163 let num_words = m_ceil_bits.div_ceil(64);
164 let num_bits = num_words * 64;
165 let k_raw = (num_bits as f64 / n) * ln_2;
167 let num_hashes = (f64_ceil_to_u64(k_raw) as u32).clamp(1, NUM_HASHES_MAX);
168 Self {
169 bits: vec![0u64; num_words as usize],
170 num_bits,
171 num_hashes,
172 }
173 }
174
175 fn from_params(num_bits: u64, num_hashes: u32, bits: Vec<u64>) -> Result<Self, BloomError> {
180 if num_bits == 0 || !num_bits.is_multiple_of(64) {
181 return Err(BloomError::BadShape(format!(
182 "num_bits {num_bits} must be a positive multiple of 64"
183 )));
184 }
185 if num_hashes == 0 || num_hashes > NUM_HASHES_MAX {
186 return Err(BloomError::BadNumHashes { got: num_hashes });
187 }
188 let expected_words = num_bits / 64;
189 if bits.len() as u64 != expected_words {
190 return Err(BloomError::BadShape(format!(
191 "bits.len() = {} doesn't match num_bits/64 = {expected_words}",
192 bits.len()
193 )));
194 }
195 Ok(Self {
196 bits,
197 num_bits,
198 num_hashes,
199 })
200 }
201
202 pub fn insert(&mut self, key: &[u8]) {
204 let (h1, h2) = derive_hash_pair(key);
205 for i in 0..self.num_hashes {
206 let bit_idx = mix(h1, h2, i, self.num_bits);
207 let word_idx = (bit_idx / 64) as usize;
208 let bit_in_word = bit_idx % 64;
209 self.bits[word_idx] |= 1u64 << bit_in_word;
210 }
211 }
212
213 #[must_use]
217 pub fn contains(&self, key: &[u8]) -> bool {
218 let (h1, h2) = derive_hash_pair(key);
219 for i in 0..self.num_hashes {
220 let bit_idx = mix(h1, h2, i, self.num_bits);
221 let word_idx = (bit_idx / 64) as usize;
222 let bit_in_word = bit_idx % 64;
223 if self.bits[word_idx] & (1u64 << bit_in_word) == 0 {
224 return false;
225 }
226 }
227 true
228 }
229
230 #[must_use]
233 pub const fn num_bits(&self) -> u64 {
234 self.num_bits
235 }
236
237 #[must_use]
239 pub const fn num_hashes(&self) -> u32 {
240 self.num_hashes
241 }
242
243 #[must_use]
246 pub fn encoded_len(&self) -> usize {
247 20 + self.bits.len() * 8
248 }
249
250 #[must_use]
254 pub fn to_bytes(&self) -> Vec<u8> {
255 let mut out = Vec::with_capacity(self.encoded_len());
256 out.extend_from_slice(&BLOOM_MAGIC.to_le_bytes());
257 let body_start = out.len();
261 out.extend_from_slice(&self.num_bits.to_le_bytes());
262 out.extend_from_slice(&self.num_hashes.to_le_bytes());
263 let crc_offset = out.len();
265 out.extend_from_slice(&0u32.to_le_bytes());
266 for word in &self.bits {
268 out.extend_from_slice(&word.to_le_bytes());
269 }
270 let body_crc = {
274 let mut to_hash = Vec::with_capacity(out.len() - crc_offset - 4 + 12);
275 to_hash.extend_from_slice(&out[body_start..crc_offset]);
276 to_hash.extend_from_slice(&out[crc_offset + 4..]);
277 crc32(&to_hash)
278 };
279 out[crc_offset..crc_offset + 4].copy_from_slice(&body_crc.to_le_bytes());
280 out
281 }
282
283 pub fn from_bytes(input: &[u8]) -> Result<Self, BloomError> {
288 const HEADER_LEN: usize = 20;
289 if input.len() < HEADER_LEN {
290 return Err(BloomError::TooShort {
291 got: input.len(),
292 need: HEADER_LEN,
293 });
294 }
295 let magic = u32::from_le_bytes([input[0], input[1], input[2], input[3]]);
296 if magic != BLOOM_MAGIC {
297 return Err(BloomError::BadMagic { got: magic });
298 }
299 let num_bits = u64::from_le_bytes([
300 input[4], input[5], input[6], input[7], input[8], input[9], input[10], input[11],
301 ]);
302 let num_hashes = u32::from_le_bytes([input[12], input[13], input[14], input[15]]);
303 let crc_stored = u32::from_le_bytes([input[16], input[17], input[18], input[19]]);
304 if num_bits == 0 || !num_bits.is_multiple_of(64) {
307 return Err(BloomError::BadShape(format!(
308 "num_bits {num_bits} must be a positive multiple of 64"
309 )));
310 }
311 let expected_words = (num_bits / 64) as usize;
312 let expected_body_bytes = expected_words * 8;
313 if input.len() != HEADER_LEN + expected_body_bytes {
314 return Err(BloomError::BadShape(format!(
315 "input is {} bytes, expected {}",
316 input.len(),
317 HEADER_LEN + expected_body_bytes
318 )));
319 }
320 let crc_computed = {
323 let mut to_hash = Vec::with_capacity(12 + expected_body_bytes);
324 to_hash.extend_from_slice(&input[4..16]); to_hash.extend_from_slice(&input[HEADER_LEN..]);
326 crc32(&to_hash)
327 };
328 if crc_computed != crc_stored {
329 return Err(BloomError::BadCrc {
330 expected: crc_stored,
331 got: crc_computed,
332 });
333 }
334 let mut bits = Vec::with_capacity(expected_words);
336 for w in 0..expected_words {
337 let off = HEADER_LEN + w * 8;
338 bits.push(u64::from_le_bytes([
339 input[off],
340 input[off + 1],
341 input[off + 2],
342 input[off + 3],
343 input[off + 4],
344 input[off + 5],
345 input[off + 6],
346 input[off + 7],
347 ]));
348 }
349 Self::from_params(num_bits, num_hashes, bits)
350 }
351}
352
353fn fnv1a_64(bytes: &[u8]) -> u64 {
356 let mut h = FNV_OFFSET_BASIS;
357 for &b in bytes {
358 h ^= u64::from(b);
359 h = h.wrapping_mul(FNV_PRIME);
360 }
361 h
362}
363
364const fn splitmix64(mut x: u64) -> u64 {
369 x = x.wrapping_add(0x9e37_79b9_7f4a_7c15);
370 x = (x ^ (x >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
371 x = (x ^ (x >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
372 x ^ (x >> 31)
373}
374
375fn derive_hash_pair(key: &[u8]) -> (u64, u64) {
376 let h1 = fnv1a_64(key);
377 let h2 = splitmix64(h1);
378 let h2 = if h2 == 0 { 0xdead_beef_dead_beef } else { h2 };
383 (h1, h2)
384}
385
386#[inline]
387fn mix(h1: u64, h2: u64, i: u32, num_bits: u64) -> u64 {
388 let combined = h1.wrapping_add((u64::from(i)).wrapping_mul(h2));
389 combined % num_bits
390}
391
392fn f64_ceil_to_u64(x: f64) -> u64 {
398 debug_assert!(x >= 0.0, "f64_ceil_to_u64: x must be >= 0");
399 let truncated = x as u64;
400 if (truncated as f64) < x {
401 truncated + 1
402 } else {
403 truncated
404 }
405}
406
407fn libm_ln(x: f64) -> f64 {
416 debug_assert!(x > 0.0, "libm_ln: x must be > 0");
417 let bits = x.to_bits();
419 let exponent_raw = ((bits >> 52) & 0x7ff) as i64;
420 let exponent = exponent_raw - 1023;
421 let mantissa_bits = (bits & 0x000f_ffff_ffff_ffff) | 0x3ff0_0000_0000_0000;
422 let mantissa = f64::from_bits(mantissa_bits);
423 use core::f64::consts::LN_2;
426 let y = mantissa - 1.0;
430 let t = y / (mantissa + 1.0);
436 let t2 = t * t;
437 let ln_mantissa = 2.0 * (t + t2 * t / 3.0 + t2 * t2 * t / 5.0 + t2 * t2 * t2 * t / 7.0);
438 (exponent as f64) * LN_2 + ln_mantissa
439}
440
441#[cfg(test)]
442mod tests {
443 use super::*;
444 use alloc::vec::Vec;
445
446 fn rng_stream(seed: u64, count: usize) -> Vec<u64> {
449 let mut s = seed;
450 let mut out = Vec::with_capacity(count);
451 for _ in 0..count {
452 s = splitmix64(s.wrapping_add(1));
453 out.push(s);
454 }
455 out
456 }
457
458 #[test]
459 fn libm_ln_matches_known_values() {
460 use core::f64::consts::{LN_2, LN_10};
462 let cases = [
463 (1.0_f64, 0.0_f64),
464 (2.0, LN_2),
465 (10.0, LN_10),
466 (0.5, -LN_2),
467 (0.01, -2.0 * LN_10),
468 ];
469 for &(x, expected) in &cases {
470 let got = libm_ln(x);
471 let err = (got - expected).abs();
472 assert!(
473 err < 1e-5,
474 "ln({x}) expected {expected}, got {got}, err {err}"
475 );
476 }
477 }
478
479 #[test]
480 fn with_target_fp_rate_sizes_match_spec() {
481 let bf = BloomFilter::with_target_fp_rate(100_000, 0.01);
484 assert_eq!(bf.num_bits() % 64, 0);
485 assert!(bf.num_bits() >= 958_506);
486 assert!(bf.num_bits() <= 958_506 + 64);
487 assert_eq!(bf.num_hashes(), 7);
488 }
489
490 #[test]
491 fn insert_then_contains_returns_true_for_inserted_keys() {
492 let mut bf = BloomFilter::with_target_fp_rate(10_000, 0.01);
493 let keys = rng_stream(0xc0ffee, 10_000);
494 for k in &keys {
495 bf.insert(&k.to_le_bytes());
496 }
497 for k in &keys {
500 assert!(
501 bf.contains(&k.to_le_bytes()),
502 "expected contains(inserted key {k}) == true"
503 );
504 }
505 }
506
507 #[test]
508 fn fuzz_oracle_fp_rate_under_target_x_1_2() {
509 const TARGET_FP: f64 = 0.01;
513 const N: usize = 100_000;
514 let mut bf = BloomFilter::with_target_fp_rate(N, TARGET_FP);
515 let inserted = rng_stream(0xfeed_beef, N);
516 for k in &inserted {
517 bf.insert(&k.to_le_bytes());
518 }
519 let probes = rng_stream(0xbeef_feed, N);
524 let inserted_set: alloc::collections::BTreeSet<u64> = inserted.iter().copied().collect();
525 let mut fp = 0u64;
526 let mut tested = 0u64;
527 for k in &probes {
528 if inserted_set.contains(k) {
529 continue;
530 }
531 tested += 1;
532 if bf.contains(&k.to_le_bytes()) {
533 fp += 1;
534 }
535 }
536 let observed = fp as f64 / tested as f64;
537 let ceiling = TARGET_FP * 1.2;
538 assert!(
539 observed <= ceiling,
540 "observed FP {observed:.4} exceeded ceiling {ceiling:.4} (target {TARGET_FP})"
541 );
542 }
543
544 #[test]
545 fn to_bytes_then_from_bytes_roundtrip() {
546 let mut bf = BloomFilter::with_target_fp_rate(1_000, 0.005);
547 let keys = rng_stream(42, 500);
548 for k in &keys {
549 bf.insert(&k.to_le_bytes());
550 }
551 let bytes = bf.to_bytes();
552 assert_eq!(bytes.len(), bf.encoded_len());
553 let parsed = BloomFilter::from_bytes(&bytes).expect("roundtrip parses");
554 assert_eq!(parsed.num_bits(), bf.num_bits());
555 assert_eq!(parsed.num_hashes(), bf.num_hashes());
556 for k in &keys {
558 assert!(parsed.contains(&k.to_le_bytes()));
559 }
560 assert_eq!(parsed.bits, bf.bits);
562 }
563
564 #[test]
565 fn from_bytes_rejects_truncated_input() {
566 let bf = BloomFilter::with_target_fp_rate(100, 0.01);
567 let bytes = bf.to_bytes();
568 let truncated = &bytes[..10];
570 match BloomFilter::from_bytes(truncated) {
571 Err(BloomError::TooShort { .. }) => {}
572 other => panic!("expected TooShort, got {other:?}"),
573 }
574 }
575
576 #[test]
577 fn from_bytes_rejects_bad_magic() {
578 let bf = BloomFilter::with_target_fp_rate(100, 0.01);
579 let mut bytes = bf.to_bytes();
580 bytes[0] ^= 0xff;
581 match BloomFilter::from_bytes(&bytes) {
582 Err(BloomError::BadMagic { .. }) => {}
583 other => panic!("expected BadMagic, got {other:?}"),
584 }
585 }
586
587 #[test]
588 fn from_bytes_rejects_bad_crc() {
589 let bf = BloomFilter::with_target_fp_rate(100, 0.01);
590 let mut bytes = bf.to_bytes();
591 bytes[25] ^= 0x01;
593 match BloomFilter::from_bytes(&bytes) {
594 Err(BloomError::BadCrc { .. }) => {}
595 other => panic!("expected BadCrc, got {other:?}"),
596 }
597 }
598
599 #[test]
600 fn from_bytes_rejects_zero_num_hashes() {
601 let num_bits: u64 = 128;
605 let num_hashes: u32 = 0;
606 let mut buf = Vec::new();
607 buf.extend_from_slice(&BLOOM_MAGIC.to_le_bytes());
608 buf.extend_from_slice(&num_bits.to_le_bytes());
609 buf.extend_from_slice(&num_hashes.to_le_bytes());
610 let crc_off = buf.len();
612 buf.extend_from_slice(&0u32.to_le_bytes());
613 for _ in 0..2 {
614 buf.extend_from_slice(&0u64.to_le_bytes());
615 }
616 let body_crc = {
617 let mut to_hash = Vec::new();
618 to_hash.extend_from_slice(&buf[4..16]);
619 to_hash.extend_from_slice(&buf[20..]);
620 crc32(&to_hash)
621 };
622 buf[crc_off..crc_off + 4].copy_from_slice(&body_crc.to_le_bytes());
623 match BloomFilter::from_bytes(&buf) {
624 Err(BloomError::BadNumHashes { got: 0 }) => {}
625 other => panic!("expected BadNumHashes, got {other:?}"),
626 }
627 }
628
629 #[test]
630 fn num_bits_is_always_64_aligned() {
631 for &(n, p) in &[
632 (1_usize, 0.5_f64),
633 (10, 0.1),
634 (1_000, 0.01),
635 (1_000_000, 0.001),
636 ] {
637 let bf = BloomFilter::with_target_fp_rate(n, p);
638 assert_eq!(bf.num_bits() % 64, 0, "n={n} p={p}");
639 assert!(bf.num_bits() >= 64);
640 }
641 }
642}