ext2/
bitmap.rs

1// Copyright 2024 The ChromiumOS Authors
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5//! Defines bitmap types.
6
7use anyhow::bail;
8use anyhow::Result;
9
10pub struct BitMap<'a> {
11    inner: &'a mut [u8],
12}
13
14impl<'a> BitMap<'a> {
15    /// Creates a new bitmap from an underlying buffer.
16    pub fn from_slice_mut(inner: &'a mut [u8]) -> Self {
17        Self { inner }
18    }
19
20    /// Sets the bit at the given index.
21    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    /// Marks the first `n` bits in the bitmap with the given value.
36    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    // Returns the number of bits in the bitmap.
51    pub fn len(&self) -> usize {
52        self.inner.len() * 8
53    }
54}
55
56// Implements test utility methods.
57#[cfg(test)]
58impl BitMap<'_> {
59    // Returns the number of bits in the bitmap that are set.
60    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        // (28 + 1) = 8 * 3 + 4. So, the first 3 bytes and 4 bits should be set.
79        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        // (42 +  1) == 8 * 5 + 3. So, 3rd bit at 6th byte should be set.
88        // Note that "+1" is needed as this is 0-indexed.
89        assert_eq!(b.as_slice(), &[0, 0, 0, 0, 0, 0b100, 0, 0, 0, 0]);
90    }
91}