diff --git a/src/lib.rs b/src/lib.rs index 3d27899..3372f6c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -551,6 +551,58 @@ impl SmallBitVec { } } + /// Counts the number of bits in the vector that are set to one/true + pub fn count_ones(&self) -> usize { + let mut len = self.len(); + if len == 0 { + return 0; + } + + if self.is_inline() { + let mask = inline_ones(len); + (self.data & mask).count_ones() as usize + } else { + let mut count = 0; + for &storage in self.buffer() { + if len >= bits_per_storage() { + count += storage.count_ones() as usize; + len -= bits_per_storage(); + } else { + let mask = (1 << len) - 1; + count += (storage & mask).count_ones() as usize; + break; + } + } + count + } + } + + /// Counts the number of bits in the vector that are set to zero/false + pub fn count_zeros(&self) -> usize { + let mut len = self.len(); + if len == 0 { + return 0; + } + + if self.is_inline() { + let mask = inline_ones(len); + (!self.data & mask).count_ones() as usize + } else { + let mut count = 0; + for &storage in self.buffer() { + if len >= bits_per_storage() { + count += storage.count_zeros() as usize; + len -= bits_per_storage(); + } else { + let mask = (1 << len) - 1; + count += (!storage & mask).count_ones() as usize; + break; + } + } + count + } + } + /// Shorten the vector, keeping the first `len` elements and dropping the rest. /// /// If `len` is greater than or equal to the vector's current length, this has no diff --git a/src/tests.rs b/src/tests.rs index 2378a5b..9c4d96d 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -365,3 +365,30 @@ fn resize() { assert_eq!(v.get(98).unwrap(), false); assert_eq!(v.get(99).unwrap(), false); } + +#[test] +fn count_ones_zeros() { + let random_bits = &[ + false, true, true, true, true, false, true, false, + true, true, false, false, false, false, true, false, + false, true, true, false, false, false, false, false, + false, true, false, false, false, false, false, true, + false, true, true, true, false, true, false, true, + true, true, false, false, true, false, false, false, + true, false, true, false, true, false, false, false, + false, true, false, true, false, false, true, true, + true, true, false, true, true, true, false, false, + false, false, true, false, true, true, false, false, + ]; + let mut v = SmallBitVec::new(); + for start in 0..random_bits.len() { + for i in 0..random_bits.len() { + let index = (start + i) % random_bits.len(); + let bit = random_bits[index]; + v.push(bit); + assert_eq!(v.iter().filter(|&b| b).count(), v.count_ones()); + assert_eq!(v.iter().filter(|&b| !b).count(), v.count_zeros()); + } + v.clear(); + } +}