bit-set/src/lib.rs

use std::ops;
use std::iter::FromIterator;
use std::default::Default;
use std::hash::{BuildHasherDefault, Hasher};
use std::collections::{HashMap, HashSet};

const BITS: u64 = 64;

struct BitHasher(u64);

impl Default for BitHasher {
    #[inline]
    fn default() -> BitHasher {
        BitHasher(0)
    }
}

impl Hasher for BitHasher {
    #[inline]
    fn finish(&self) -> u64 {
        self.0
    }

    #[inline]
    fn write(&mut self, _: &[u8]) {
        panic!("this hasher only works with u64");
    }

    #[inline]
    fn write_u64(&mut self, i: u64) {
        *self = BitHasher(i);
    }
}

type BitBuildHasher = BuildHasherDefault<BitHasher>;
type BitHashMap<K, V> = HashMap<K, V, BitBuildHasher>;
// type BitHashSet<T> = HashSet<T, BitBuildHasher>;

type Block = u64;
type Storage = BitHashMap<u64, Block>;

#[inline]
fn block_bit(x: u64, d: u64) -> (u64, u64) {
    (x / d, x % d)
}

#[derive(Debug)]
pub struct BitSet {
    blocks: Storage,
    nbits: usize,
}

impl BitSet {
    #[inline]
    pub fn new() -> BitSet {
        BitSet {
            blocks: Storage::default(),
            nbits: 0,
        }
    }

    #[inline]
    pub fn with_capacity(nbits: usize) -> BitSet {
        let (mut block, bit) = block_bit(nbits as u64, BITS);
        block += (bit > 0) as u64;

        BitSet {
            blocks: Storage::with_capacity_and_hasher(block as usize, Default::default()),
            nbits: 0,
        }
    }

    #[inline]
    pub fn capacity(&self) -> usize {
        self.blocks.len() * BITS as usize
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.nbits
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.nbits == 0
    }

    #[inline]
    pub fn insert(&mut self, value: u64) -> bool {
        let (block, bit) = block_bit(value, BITS);
        let block = self.blocks.entry(block).or_insert(0);

        let n = 1 << bit;

        if (*block & n) == 0 {
            *block |= n;
            self.nbits += 1;

            true
        } else {
            false
        }
    }

    #[inline]
    pub fn remove(&mut self, value: u64) -> bool {
        let (block, bit) = block_bit(value, BITS);
        let block = self.blocks.entry(block).or_insert(0);

        let n = 1 << bit;

        if (*block & n) != 0 {
            *block &= !n;
            self.nbits -= 1;

            true
        } else {
            false
        }
    }

    #[inline]
    pub fn contains(&self, value: u64) -> bool {
        let (block, bit) = block_bit(value, BITS);

        match self.blocks.get(&block) {
            None => false,
            Some(block) => (block & (1 << bit)) != 0,
        }
    }

    #[inline]
    pub fn extend_from_bitset(&mut self, other: &Self) {
        for (key, value) in &other.blocks {
            *self.blocks.entry(*key).or_insert(0) |= value;
        }

        self.nbits = self.blocks
            .values()
            .map(|block| block.count_ones() as usize)
            .sum();
    }
}

impl Default for BitSet {
    #[inline]
    fn default() -> BitSet {
        BitSet::new()
    }
}

impl FromIterator<u64> for BitSet {
    #[inline]
    fn from_iter<I: IntoIterator<Item = u64>>(iter: I) -> BitSet {
        let mut set = BitSet::new();
        set.extend(iter);
        set
    }
}

impl Extend<u64> for BitSet {
    #[inline]
    fn extend<I: IntoIterator<Item = u64>>(&mut self, iter: I) {
        for i in iter {
            self.insert(i);
        }
    }
}

impl<'a> Extend<&'a u64> for BitSet {
    #[inline]
    fn extend<I: IntoIterator<Item = &'a u64>>(&mut self, iter: I) {
        for i in iter {
            self.insert(*i);
        }
    }
}

impl ops::BitOr for BitSet {
    type Output = Self;

    #[inline]
    fn bitor(self, rhs: Self) -> Self {
        let mut blocks = self.blocks.clone();

        for (key, value) in &rhs.blocks {
            *blocks.entry(*key).or_insert(0) |= value;
        }

        let nbits = blocks
            .values()
            .map(|block| block.count_ones() as usize)
            .sum();

        BitSet {
            blocks: blocks,
            nbits: nbits,
        }
    }
}

impl<'a> ops::BitOr<&'a Self> for BitSet {
    type Output = Self;

    #[inline]
    fn bitor(self, rhs: &'a Self) -> Self {
        let mut blocks = self.blocks.clone();

        for (key, value) in &rhs.blocks {
            *blocks.entry(*key).or_insert(0) |= value;
        }

        let nbits = blocks
            .values()
            .map(|block| block.count_ones() as usize)
            .sum();

        BitSet {
            blocks: blocks,
            nbits: nbits,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn insert() {
        let mut set = BitSet::with_capacity(10);

        assert_eq!(set.contains(0), false);
        assert_eq!(set.contains(10), false);

        set.insert(0);
        set.insert(10);

        assert_eq!(set.contains(0), true);
        assert_eq!(set.contains(10), true);

        assert_eq!(set.contains(100), false);

        set.insert(100);

        assert_eq!(set.contains(100), true);
    }

    #[test]
    fn from_iter() {
        let set = [1, 2, 3, 10, 100].iter().cloned().collect::<BitSet>();

        assert_eq!(set.len(), 5);

        assert_eq!(set.contains(1), true);
        assert_eq!(set.contains(2), true);
        assert_eq!(set.contains(3), true);
        assert_eq!(set.contains(10), true);
        assert_eq!(set.contains(100), true);
    }

    #[test]
    fn bitor() {
        let set_a = [1, 2, 3].iter().cloned().collect::<BitSet>();
        let set_b = [3, 4, 5].iter().cloned().collect::<BitSet>();

        let set = set_a | set_b;

        assert_eq!(set.len(), 5);

        assert_eq!(set.contains(1), true);
        assert_eq!(set.contains(2), true);
        assert_eq!(set.contains(3), true);
        assert_eq!(set.contains(4), true);
        assert_eq!(set.contains(5), true);
    }
}