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 7 use anyhow::bail; 8 use anyhow::Result; 9 10 pub struct BitMap<'a> { 11 inner: &'a mut [u8], 12 } 13 14 impl<'a> BitMap<'a> { 15 /// Creates a new bitmap from an underlying buffer. from_slice_mut(inner: &'a mut [u8]) -> Self16 pub fn from_slice_mut(inner: &'a mut [u8]) -> Self { 17 Self { inner } 18 } 19 20 /// Sets the bit at the given index. set(&mut self, index: usize, value: bool) -> Result<()>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. mark_first_elems(&mut self, n: usize, value: bool)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. len(&self) -> usize51 pub fn len(&self) -> usize { 52 self.inner.len() * 8 53 } 54 } 55 56 // Implements test utility methods. 57 #[cfg(test)] 58 impl<'a> BitMap<'a> { 59 // Returns the number of bits in the bitmap that are set. count_zeros(&self) -> usize60 pub fn count_zeros(&self) -> usize { 61 self.inner.iter().map(|b| b.count_zeros() as usize).sum() 62 } 63 as_slice(&self) -> &[u8]64 pub fn as_slice(&self) -> &[u8] { 65 self.inner 66 } 67 } 68 69 #[cfg(test)] 70 mod tests { 71 use super::*; 72 73 #[test] test_mark_first_elems()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] test_set()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 } 92