Documentation

extern crate bytes;

use bytes::{BufMut, BytesMut, Bytes, Buf};

fn get_u32_little_end(arr: &[u8]) -> u32 {
    ((arr[0] as u32) <<  0) +
        ((arr[1] as u32) <<  8) +
        ((arr[2] as u32) << 16) +
        ((arr[3] as u32) << 24)
}

fn hash(data: &Bytes, seed: u32) -> u32 {
    let m = 0xc6a4a793 as u32;
    let r = 24 as u32;
    let mut h = seed ^ (data.len() as u64 * m as u64) as u32;
    let n = data.len() - data.len() % 4;
    let mut i = 0;

    while i < n {
        h = (h as u64 + get_u32_little_end(&data[i..]) as u64) as u32;
        h = (h as u64 * m as u64) as u32;
        h ^= (h >> 16);
        i += 4;
    }

    let flag = data.len() - i;
    if flag == 3 {
        h += (data.len() as u32) << 16;
    } else if flag == 2 {
        h += (data.len() as u32) << 8;
    } else if flag == 1 {
        h += data[i] as u32;
        h = (h as u64 * m as u64) as u32;
        h ^= (h >> r);
    }

    return h
}

fn bloom_hash(data: &Bytes) -> u32 {
    hash(data, 0xbc9f1d34)
}

pub struct BloomFilter {
    bits_per_key: usize,
    k: u8,
    key_hashes: Vec<u32>,
}

impl BloomFilter {

    pub fn new(bits_per_key: usize) -> Self {
        let mut k = (bits_per_key as f64 * 0.69) as u8;
        if k < 1 { k = 1; }
        if k > 30 { k = 30; }
        BloomFilter { bits_per_key, k, key_hashes: vec![] }
    }

    pub fn contains(&self, filter: &Bytes, key: &Bytes) -> bool {
        let n_bytes = filter.len() - 1;
        if n_bytes < 1 {
            return false
        }
        let n_bits = (n_bytes * 8) as u32;

        let k = filter[n_bytes];
        if k > 30 {
            return true
        }

        let mut kh = bloom_hash(key);
        let delta = (kh >> 17) | (kh << 15);
        for _ in 0..k {
            let bitpos = (kh % n_bits) as usize;
            if filter[bitpos/8] as u32 & (1 << (bitpos % 8)) == 0 {
                return false
            }
            kh = (kh as u64 + delta as u64) as u32;
        }

        return true
    }

    pub fn add(&mut self, key: &Bytes) {
        self.key_hashes.push(bloom_hash(key))
    }

    pub fn generate(&mut self) -> Bytes {
        let mut n_bits = (self.key_hashes.len() * self.bits_per_key) as u32;
        if n_bits < 64 {
            n_bits = 64;
        }
        let n_bytes = (n_bits + 7) / 8;
        n_bits = n_bytes * 8;

        let mut dest = BytesMut::new();
        dest.resize(n_bytes as usize + 1, 0);
        dest[n_bytes as usize] = self.k;
        for v in &self.key_hashes {
            let mut kh = v.clone();
            let delta = (kh >> 17) | (kh << 15);
            for _ in 0..self.k {
                let bitpos = (kh % n_bits) as usize;
                dest[bitpos/8] |= (1 << (bitpos % 8));
                kh = (kh as u64 + delta as u64) as u32;
            }
        }
        self.key_hashes.clear();

        dest.freeze()
    }

}

#[cfg(test)]
mod tests {
    use bytes::{Bytes, BytesMut, BufMut, Buf};
    use crate::BloomFilter;

    fn num_to_bytes(num: u32) -> Bytes {
        let mut bs = BytesMut::new();
        bs.put_u32_le(num);
        bs.freeze()
    }

    #[test]
    fn it_works() {
        let mut bf = BloomFilter::new(10);
        let n = 10000;
        for i in 0..n {
            bf.add(&num_to_bytes(i))
        }
        let filter = bf.generate();

        for i in 0..n {
            if !bf.contains(&filter, &num_to_bytes(i)) {
                panic!(format!("Error in {}", i))
            }
        }

        let mut rate: f32 = 0.0;
        for i in 0..n {
            if bf.contains(&filter, &num_to_bytes(i + n + 1)) {
                rate += 1.0;
            }
        }

        rate /= n as f32;
        if rate > 0.02 {
            panic!(format!("False positive rate is more than 2%%, got {}, at len {}", rate, n))
        } else {
            println!("False positive rate is {}", rate)
        }
    }
}