algo_rs/
bitset.rs

1use std::ops::Index;
2
3pub struct BitSet {
4    size: usize,
5    inner: Vec<u64>,
6}
7
8#[inline(always)]
9fn mask(bit: usize) -> u64 {
10    1 << (bit % 64)
11}
12
13impl BitSet {
14    pub fn new(size: usize) -> Self {
15        Self {
16            size,
17            inner: vec![0; (size + 63) / 64],
18        }
19    }
20
21    pub fn get(&self, bit: usize) -> Option<bool> {
22        if bit >= self.size {
23            return None;
24        }
25
26        let mask = mask(bit);
27
28        Some(self.inner[bit / 64] & mask == mask)
29    }
30
31    pub fn set(&mut self, bit: usize, value: bool) {
32        if bit >= self.size {
33            panic!("{} index is out of bitset bounds", bit)
34        }
35
36        if value {
37            self.inner[bit / 64] |= mask(bit);
38        } else {
39            self.inner[bit / 64] &= !mask(bit);
40        }
41    }
42}
43
44impl Index<usize> for BitSet {
45    type Output = bool;
46
47    fn index(&self, bit: usize) -> &Self::Output {
48        if self.get(bit).unwrap() {
49            &true
50        } else {
51            &false
52        }
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use crate::bitset::BitSet;
59
60    #[test]
61    fn test_bit_set_get_set() {
62        let mut bit_set = BitSet::new(1000);
63        for i in 1..1000 {
64            if i & (i - 1) == 0 {
65                bit_set.set(i, true);
66            }
67        }
68
69        for i in 1..1000 {
70            assert_eq!(i & (i - 1) == 0, bit_set.get(i).unwrap());
71        }
72
73        assert_eq!(None, bit_set.get(1000));
74    }
75
76    #[test]
77    fn test_bit_set_index() {
78        let mut bit_set = BitSet::new(1000);
79        for i in 1..1000 {
80            if i & (i - 1) == 0 {
81                bit_set.set(i, true);
82            }
83        }
84
85        for i in 1..1000 {
86            assert_eq!(i & (i - 1) == 0, bit_set[i]);
87        }
88    }
89
90    #[test]
91    #[should_panic(expected = "1001 index is out of bitset bounds")]
92    fn test_fixed_stack_full_stack() {
93        let mut bit_set = BitSet::new(1000);
94        bit_set.set(1001, false);
95    }
96}