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}