1use anyhow::bail;
8use anyhow::Result;
9
10pub struct BitMap<'a> {
11 inner: &'a mut [u8],
12}
13
14impl<'a> BitMap<'a> {
15 pub fn from_slice_mut(inner: &'a mut [u8]) -> Self {
17 Self { inner }
18 }
19
20 pub fn set(&mut self, index: usize, value: bool) -> Result<()> {
22 let i = index / 8;
23 let j = index % 8;
24 if i >= self.inner.len() {
25 bail!("index out of range");
26 }
27 if value {
28 self.inner[i] |= 1 << j;
29 } else {
30 self.inner[i] &= !(1 << j);
31 }
32 Ok(())
33 }
34
35 pub fn mark_first_elems(&mut self, n: usize, value: bool) {
37 self.inner
38 .iter_mut()
39 .take(n / 8)
40 .for_each(|v| *v = if value { 0xff } else { 0 });
41 if n % 8 != 0 {
42 if value {
43 self.inner[n / 8] |= 0xff >> (8 - n % 8);
44 } else {
45 self.inner[n / 8] &= !(0xff >> (8 - n % 8));
46 }
47 }
48 }
49
50 pub fn len(&self) -> usize {
52 self.inner.len() * 8
53 }
54}
55
56#[cfg(test)]
58impl BitMap<'_> {
59 pub fn count_zeros(&self) -> usize {
61 self.inner.iter().map(|b| b.count_zeros() as usize).sum()
62 }
63
64 pub fn as_slice(&self) -> &[u8] {
65 self.inner
66 }
67}
68
69#[cfg(test)]
70mod tests {
71 use super::*;
72
73 #[test]
74 fn test_mark_first_elems() {
75 let mut s = [0; 10];
76 let mut b = BitMap::from_slice_mut(&mut s);
77 b.mark_first_elems(28, true);
78 assert_eq!(b.as_slice(), &[0xff, 0xff, 0xff, 0b1111, 0, 0, 0, 0, 0, 0]);
80 }
81
82 #[test]
83 fn test_set() {
84 let mut s = [0; 10];
85 let mut b = BitMap::from_slice_mut(&mut s);
86 b.set(42, true).unwrap();
87 assert_eq!(b.as_slice(), &[0, 0, 0, 0, 0, 0b100, 0, 0, 0, 0]);
90 }
91}