src/bitset.rs


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

#[derive(Clone, PartialEq, Eq)]
pub struct BitSet {
    data: Vec<u32>,
    size: usize,
}

const BITS: usize = 32;

impl BitSet {
    pub fn new(size: usize) -> Self {
        Self {
            size,
            data: vec![0; size.div_ceil(32)],
        }
    }

    pub fn set(&mut self, bit: usize, val: bool) {
        assert!(bit < self.size);
        let idx = bit / BITS;
        let bit = bit % BITS;

        if val {
            self.data[idx] |= 1 << bit;
        } else {
            self.data[idx] &= !(1 << bit);
        }
    }

    pub fn get(&self, bit: usize) -> bool {
        assert!(bit < self.size);
        let idx = bit / BITS;
        let bit = bit % BITS;

        (self.data[idx] & (1 << bit)) != 0
    }

    pub fn set_all(&mut self, val: bool) {
        let val = if val { !0 } else { 0 };
        for d in self.data.iter_mut() {
            *d = val;
        }
    }
}