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
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
#![feature(const_generics)]
#![feature(const_evaluatable_checked)]

const USIZE_BITS: usize = usize::BITS as usize;

#[derive(Debug)]
pub enum StackBitSetError {
    IndexOutOfBounds,
}

pub const fn usize_count(n: usize) -> usize {
    (n / USIZE_BITS) + if n % USIZE_BITS == 0 { 0 } else { 1 }
}

#[derive(Clone, Copy, Debug)]
pub struct StackBitSet<const N: usize>
where
    [(); usize_count(N)]: Sized,
{
    data: [usize; usize_count(N)],
}

impl<const N: usize> StackBitSet<N>
where
    [(); usize_count(N)]: Sized,
{
    pub fn new() -> Self {
        StackBitSet {
            data: [0usize; usize_count(N)],
        }
    }
    pub fn get(&self, idx: usize) -> Result<bool, StackBitSetError> {
        if idx < N {
            Ok(self.get_unchecked(idx))
        } else {
            Err(StackBitSetError::IndexOutOfBounds)
        }
    }
    fn get_unchecked(&self, idx: usize) -> bool {
        let chunk = self.data[idx / USIZE_BITS];
        // ! Make sure this is the right cast
        chunk & (1 << (idx % USIZE_BITS)) != 0
    }
    pub fn set(&mut self, idx: usize) -> Result<(), StackBitSetError> {
        if idx < N {
            Ok(self.set_unchecked(idx))
        } else {
            Err(StackBitSetError::IndexOutOfBounds)
        }
    }
    fn set_unchecked(&mut self, idx: usize) {
        let chunk = unsafe { self.data.get_unchecked_mut(idx / USIZE_BITS) };
        *chunk = *chunk | (1 << (idx % USIZE_BITS))
    }
    pub fn reset(&mut self, idx: usize) -> Result<(), StackBitSetError> {
        if idx < N {
            Ok(self.reset_unchecked(idx))
        } else {
            Err(StackBitSetError::IndexOutOfBounds)
        }
    }
    fn reset_unchecked(&mut self, idx: usize) {
        let chunk = unsafe { self.data.get_unchecked_mut(idx / USIZE_BITS) };
        *chunk = *chunk & !(1 << (idx % USIZE_BITS))
    }
}

#[cfg(test)]
mod tests {
    use crate::StackBitSet;
    #[test]
    fn bitset_create() {
        let _a: StackBitSet<42> = StackBitSet::new();
    }

    #[test]
    fn set_reset_bit() {
        let mut a: StackBitSet<42> = StackBitSet::new();
        assert!(!a.get(12).unwrap());
        a.set(12).unwrap();
        assert!(a.get(12).unwrap());
        a.reset(12).unwrap();
        assert!(!a.get(12).unwrap());
    }
}