rustolio_utils/

bloom_filter.rs

//
// SPDX-License-Identifier: MPL-2.0
//
// Copyright (c) 2026 Tobias Binnewies. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//

#![allow(deprecated)]

use std::hash::{Hash, Hasher as _, SipHasher};

use crate::{bit_vec::BitVector, crypto::rand, prelude::*};

#[derive(Debug, Clone)]
pub struct BloomFilter {
    bits: BitVector,
    k_num: u32,
    sips: [SipHasher; 2],
    seed: [u64; 4],
}

impl PartialEq for BloomFilter {
    fn eq(&self, other: &Self) -> bool {
        self.bits == other.bits && self.k_num == other.k_num && self.seed == other.seed
    }
}
impl Eq for BloomFilter {}

/// m=−n⋅ln(p)​ / (ln2)^2
/// m = number of bits needed
/// n = expected number of items
/// p = desired false positive probability
fn calc_filter_size(n: usize, p: f64) -> usize {
    assert!(n > 0);
    assert!(p > 0.0 && p < 1.0);
    let log2_2 = std::f64::consts::LN_2 * std::f64::consts::LN_2;
    -((n as f64) * f64::ln(p) / log2_2).ceil() as usize
}

/// k= nm​ ln2
/// m = number of bits in the bitmap (bitmap_bits)
/// n = expected number of items (item_count)
/// k = optimal number of hash functions
fn optimal_k_num(m: u64, n: usize) -> u32 {
    let m = m as f64;
    let n = n as f64;
    let k_num = (m / n * std::f64::consts::LN_2).round() as u32;
    std::cmp::max(k_num, 1)
}

impl BloomFilter {
    pub fn new(item_count: usize, p: f64) -> rand::Result<Self> {
        let size = calc_filter_size(item_count, p);
        let k_num = optimal_k_num(size as u64, item_count);

        let s1 = rand::u64()?;
        let s2 = rand::u64()?;
        let s3 = rand::u64()?;
        let s4 = rand::u64()?;
        let sip1 = SipHasher::new_with_keys(s1, s2);
        let sip2 = SipHasher::new_with_keys(s3, s4);

        Ok(Self {
            bits: BitVector::new(size),
            k_num,
            sips: [sip1, sip2],
            seed: [s1, s2, s3, s4],
        })
    }

    pub fn set(&mut self, item: &impl Hash) {
        let mut hashes = [0; 2];
        for k_i in 0..self.k_num {
            let bit_offset = self.bloom_hash(&mut hashes, item, k_i) as usize % self.bits.len();
            self.bits.set(bit_offset, true);
        }
    }

    pub fn check(&self, item: &impl Hash) -> bool {
        let mut hashes = [0u64, 0u64];
        for k_i in 0..self.k_num {
            let bit_offset = self.bloom_hash(&mut hashes, item, k_i) as usize % self.bits.len();
            if !self.bits.get(bit_offset) {
                return false;
            }
        }
        true
    }

    pub fn check_and_set(&mut self, item: &impl Hash) -> bool {
        let mut hashes = [0u64, 0u64];
        let mut found = true;
        for k_i in 0..self.k_num {
            let bit_offset = self.bloom_hash(&mut hashes, item, k_i) as usize % self.bits.len();
            if !self.bits.get(bit_offset) {
                found = false;
                self.bits.set(bit_offset, true);
            }
        }
        found
    }

    fn bloom_hash(&self, hashes: &mut [u64; 2], item: &impl Hash, k_i: u32) -> u64 {
        if k_i > 1 {
            return (hashes[0]).wrapping_add((k_i as u64).wrapping_mul(hashes[1]))
                % 0xFFFF_FFFF_FFFF_FFC5u64;
            //largest u64 prime
        }
        let sip = &mut self.sips[k_i as usize].clone();
        item.hash(sip);
        let hash = sip.finish();
        hashes[k_i as usize] = hash;
        hash
    }
}

impl Encode for BloomFilter {
    #[async_impl]
    async fn encode_async(
        &self,
        writer: &mut impl __utils_macros::AsyncEncoder,
    ) -> crate::Result<()> {
        self.bits.encode_async(writer).await?;
        self.k_num.encode_async(writer).await?;
        self.seed.encode_async(writer).await?;
        Ok(())
    }

    fn encode_size(&self) -> usize {
        self.bits.encode_size() + self.k_num.encode_size() + self.seed.encode_size()
    }
}

impl Decode for BloomFilter {
    #[async_impl]
    async fn decode_async(reader: &mut impl __utils_macros::AsyncDecoder) -> crate::Result<Self> {
        let bits = BitVector::decode_async(reader).await?;
        let k_num = u32::decode_async(reader).await?;
        let seed = <[u64; 4]>::decode_async(reader).await?;
        let sip1 = SipHasher::new_with_keys(seed[0], seed[1]);
        let sip2 = SipHasher::new_with_keys(seed[2], seed[3]);
        Ok(Self {
            bits,
            k_num,
            sips: [sip1, sip2],
            seed,
        })
    }
}

#[cfg(test)]
mod tests {
    use crate::bytes::encoding::{decode_from_bytes, encode_to_bytes};

    use super::*;

    #[test]
    fn test_bloom_basic() {
        let i0 = "test";
        let i1 = "test1";
        let i2 = "test2";

        let mut bloom = BloomFilter::new(99, 0.001).unwrap();
        assert_eq!(bloom.k_num, 10);
        assert_eq!(bloom.bits.len(), 1423);

        assert!(!bloom.check(&i0));
        assert!(!bloom.check(&i1));
        assert!(!bloom.check(&i2));

        bloom.set(&i0);
        assert!(bloom.check(&i0));
        assert!(!bloom.check(&i1));
        assert!(!bloom.check(&i2));

        bloom.set(&i1);
        assert!(bloom.check(&i0));
        assert!(bloom.check(&i1));
        assert!(!bloom.check(&i2));

        bloom.set(&i2);
        assert!(bloom.check(&i0));
        assert!(bloom.check(&i1));
        assert!(bloom.check(&i2));
    }

    #[test]
    fn test_bloom_encoding() {
        let i0 = 123;
        let i1 = 456;
        let i2 = 789;

        let mut bloom = BloomFilter::new(9999, 0.0001).unwrap();
        assert_eq!(bloom.k_num, 13);
        assert_eq!(bloom.bits.len(), 191681);

        assert!(!bloom.check(&i0));
        assert!(!bloom.check(&i1));
        assert!(!bloom.check(&i2));

        bloom.set(&i0);
        assert!(bloom.check(&i0));
        assert!(!bloom.check(&i1));
        assert!(!bloom.check(&i2));

        let encoded = encode_to_bytes(&bloom).unwrap();
        let mut bloom: BloomFilter = decode_from_bytes(encoded).unwrap();

        assert!(bloom.check(&i0));
        assert!(!bloom.check(&i1));
        assert!(!bloom.check(&i2));

        bloom.set(&i1);
        assert!(bloom.check(&i0));
        assert!(bloom.check(&i1));
        assert!(!bloom.check(&i2));
    }
}