Optimized .count() for the iterators.

Added a private count_set_bits() method that counts set bits in O(number of words)
Fixed the deprecated rng methods in the benchmark
Added tests and benchmarks for the new .count() functionality

(softdevteam#87)

Author: eadf

benches/vob.rs

@@ -65,9 +65,9 @@ fn split_off(c: &mut Criterion) {
 fn xor(c: &mut Criterion) {
     let mut v1 = Vob::with_capacity(N);
     let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
-    v1.extend((0..N).map(|_| rng.gen::<bool>()));
+    v1.extend((0..N).map(|_| rng.random::<bool>()));
     let mut v2 = Vob::with_capacity(N);
-    v2.extend((0..N).map(|_| rng.gen::<bool>()));
+    v2.extend((0..N).map(|_| rng.random::<bool>()));
 
     c.bench_function("xor", |b| {
         b.iter(|| {
@@ -79,9 +79,9 @@ fn xor(c: &mut Criterion) {
 fn or(c: &mut Criterion) {
     let mut v1 = Vob::with_capacity(N);
     let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
-    v1.extend((0..N).map(|_| rng.gen::<bool>()));
+    v1.extend((0..N).map(|_| rng.random::<bool>()));
     let mut v2 = Vob::with_capacity(N);
-    v2.extend((0..N).map(|_| rng.gen::<bool>()));
+    v2.extend((0..N).map(|_| rng.random::<bool>()));
 
     c.bench_function("or", |b| {
         b.iter(|| {
@@ -93,9 +93,9 @@ fn or(c: &mut Criterion) {
 fn and(c: &mut Criterion) {
     let mut v1 = Vob::with_capacity(N);
     let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
-    v1.extend((0..N).map(|_| rng.gen::<bool>()));
+    v1.extend((0..N).map(|_| rng.random::<bool>()));
     let mut v2 = Vob::with_capacity(N);
-    v2.extend((0..N).map(|_| rng.gen::<bool>()));
+    v2.extend((0..N).map(|_| rng.random::<bool>()));
 
     c.bench_function("and", |b| {
         b.iter(|| {
@@ -105,38 +105,88 @@ fn and(c: &mut Criterion) {
 }
 
 fn from_bytes(c: &mut Criterion) {
-    let mut rng = rand::thread_rng();
+    let mut rng = rand::rng();
     let mut v1 = [0u8; 1024];
     rng.fill(&mut v1);
     c.bench_function("from_bytes", |b| b.iter(|| Vob::from_bytes(&v1)));
 }
 
+fn iter_bits(c: &mut Criterion) {
+    let mut a = Vob::with_capacity(N);
+    let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
+    a.extend((0..N).map(|_| rng.random::<bool>()));
+    c.bench_function("iter_bits", |b| {
+        b.iter(|| a.iter().filter(|_| true).count())
+    });
+}
+
+// this benchmark can later be removed, since it is only benchmarking std::Rage.count() as a
+// comparison with the old way of doing it. (iter_bits() above)
+fn count_bits(c: &mut Criterion) {
+    let mut a = Vob::with_capacity(N);
+    let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
+    a.extend((0..N).map(|_| rng.random::<bool>()));
+    c.bench_function("count_bits", |b| b.iter(|| a.iter().count()));
+}
+
 fn iter_set_bits(c: &mut Criterion) {
     let mut a = Vob::with_capacity(N);
     let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
-    a.extend((0..N).map(|_| rng.gen::<bool>()));
-    c.bench_function("iter_set_bits", |b| b.iter(|| a.iter_set_bits(..).count()));
+    a.extend((0..N).map(|_| rng.random::<bool>()));
+    c.bench_function("iter_set_bits", |b| {
+        b.iter(|| a.iter_set_bits(..).filter(|_| true).count())
+    });
+}
+
+fn count_set_bits(c: &mut Criterion) {
+    let mut a = Vob::with_capacity(N);
+    let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
+    a.extend((0..N).map(|_| rng.random::<bool>()));
+    c.bench_function("count_set_bits", |b| b.iter(|| a.iter_set_bits(..).count()));
 }
 
 fn iter_set_bits_u8(c: &mut Criterion) {
     let mut a = Vob::<u8>::new_with_storage_type(N);
     let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
-    a.extend((0..N).map(|_| rng.gen::<bool>()));
+    a.extend((0..N).map(|_| rng.random::<bool>()));
     c.bench_function("iter_set_bits_u8", |b| {
+        b.iter(|| a.iter_set_bits(..).filter(|_| true).count())
+    });
+}
+
+fn count_set_bits_u8(c: &mut Criterion) {
+    let mut a = Vob::<u8>::new_with_storage_type(N);
+    let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
+    a.extend((0..N).map(|_| rng.random::<bool>()));
+    c.bench_function("count_set_bits_u8", |b| {
         b.iter(|| a.iter_set_bits(..).count())
     });
 }
 
 fn iter_all_set_bits(c: &mut Criterion) {
     let a = Vob::from_elem(true, N);
     c.bench_function("iter_all_set_bits", |b| {
+        b.iter(|| a.iter_set_bits(..).filter(|_| true).count())
+    });
+}
+
+fn count_all_set_bits(c: &mut Criterion) {
+    let a = Vob::from_elem(true, N);
+    c.bench_function("count_all_set_bits", |b| {
         b.iter(|| a.iter_set_bits(..).count())
     });
 }
 
 fn iter_all_unset_bits(c: &mut Criterion) {
     let a = Vob::from_elem(true, N);
     c.bench_function("iter_all_unset_bits", |b| {
+        b.iter(|| a.iter_unset_bits(..).filter(|_| true).count())
+    });
+}
+
+fn count_all_unset_bits(c: &mut Criterion) {
+    let a = Vob::from_elem(true, N);
+    c.bench_function("count_all_unset_bits", |b| {
         b.iter(|| a.iter_unset_bits(..).count())
     });
 }
@@ -152,9 +202,15 @@ criterion_group!(
     or,
     and,
     from_bytes,
+    iter_bits,
+    count_bits,
     iter_set_bits,
+    count_set_bits,
     iter_set_bits_u8,
+    count_set_bits_u8,
     iter_all_set_bits,
-    iter_all_unset_bits
+    count_all_set_bits,
+    iter_all_unset_bits,
+    count_all_unset_bits
 );
 criterion_main!(benches);

src/lib.rs

@@ -5,6 +5,7 @@
 
 use std::{
     cmp::{min, PartialEq},
+    convert::TryFrom,
     fmt::{self, Debug},
     hash::{Hash, Hasher},
     iter::{DoubleEndedIterator, FromIterator, FusedIterator},
@@ -585,6 +586,64 @@ impl<T: Debug + PrimInt> Vob<T> {
         }
     }
 
+    /// Counts the number of set bits.
+    /// This method assumes the range is processed with process_range()
+    fn count_set_bits(&self, range: Range<usize>) -> usize {
+        // Early return for empty ranges
+        if range.is_empty() {
+            return 0;
+        }
+        let start_off = block_offset::<T>(range.start);
+        debug_assert!(
+            start_off < self.len(),
+            "start_off {} >= self.len {}",
+            start_off,
+            self.len()
+        );
+
+        // this -1 arithmetic is safe since we already tested for range.start & range.end equality
+        let end_off = blocks_required::<T>(range.end) - 1;
+
+        if start_off == end_off {
+            // Range entirely within one word
+            let b = self.vec[start_off];
+            let start_bit = range.start % bits_per_block::<T>();
+            let end_bit = range.end % bits_per_block::<T>();
+
+            // Remove bits before start_bit and bits after end_bit
+            let count = if end_bit == 0 {
+                // end_bit = 0 means we want everything from start_bit to end of word
+                // After the right shift, we have what we want
+                b >> start_bit
+            } else {
+                // We want bits from start_bit to end_bit
+                // After right shift, we need to remove the high bits
+                (b >> start_bit) << (start_bit + bits_per_block::<T>() - end_bit)
+            }
+            .count_ones();
+            return usize::try_from(count).unwrap();
+        }
+
+        // First word: shift out bits before start_bit
+        let start_bit = range.start % bits_per_block::<T>();
+        let mut count = usize::try_from((self.vec[start_off] >> start_bit).count_ones()).unwrap();
+
+        // Middle words
+        for word_idx in (start_off + 1)..end_off {
+            count += usize::try_from(self.vec[word_idx].count_ones()).unwrap();
+        }
+
+        // Last word: shift out bits after end_bit
+        let end_bit = range.end % bits_per_block::<T>();
+        let count_ones = if end_bit == 0 {
+            // end_bit = 0 means we want to count the entire end_off word
+            self.vec[end_off].count_ones()
+        } else {
+            (self.vec[end_off] << (bits_per_block::<T>() - end_bit)).count_ones()
+        };
+        count + usize::try_from(count_ones).unwrap()
+    }
+
     /// Returns an iterator which efficiently produces the index of each unset bit in the specified
     /// range. Assuming appropriate support from your CPU, this is much more efficient than
     /// checking each bit individually.
@@ -1082,6 +1141,10 @@ impl<T: Debug + PrimInt> Iterator for Iter<'_, T> {
     fn size_hint(&self) -> (usize, Option<usize>) {
         self.range.size_hint()
     }
+
+    fn count(self) -> usize {
+        self.range.count()
+    }
 }
 
 impl<T: Debug + PrimInt> DoubleEndedIterator for Iter<'_, T> {
@@ -1148,6 +1211,10 @@ impl<T: Debug + PrimInt> Iterator for IterSetBits<'_, T> {
     fn size_hint(&self) -> (usize, Option<usize>) {
         self.range.size_hint()
     }
+
+    fn count(self) -> usize {
+        self.vob.count_set_bits(self.range)
+    }
 }
 
 impl<T: Debug + PrimInt> DoubleEndedIterator for IterSetBits<'_, T> {
@@ -1228,6 +1295,12 @@ impl<T: Debug + PrimInt> Iterator for IterUnsetBits<'_, T> {
     fn size_hint(&self) -> (usize, Option<usize>) {
         self.range.size_hint()
     }
+
+    fn count(self) -> usize {
+        // This arithmetic is safe because (self.range.end - self.range.start) is the total number of bits,
+        // and self.vob.count_set_bits() always returns a value less than or equal to that.
+        (self.range.end - self.range.start) - self.vob.count_set_bits(self.range)
+    }
 }
 
 impl<T: Debug + PrimInt> DoubleEndedIterator for IterUnsetBits<'_, T> {
@@ -1300,6 +1373,10 @@ impl<T: Debug + PrimInt> Iterator for StorageIter<'_, T> {
     fn size_hint(&self) -> (usize, Option<usize>) {
         self.iter.size_hint()
     }
+
+    fn count(self) -> usize {
+        self.iter.count()
+    }
 }
 
 #[inline(always)]
@@ -1974,6 +2051,27 @@ mod tests {
         for _ in 0..len {
             vob.push(rng.random());
         }
+        // these tests can later be dialed down, as they noticeable slow down every random vob test.
+        assert_eq!(
+            vob.iter_set_bits(..).count(),
+            vob.iter_set_bits(..).filter(|_| true).count()
+        );
+        assert_eq!(
+            vob.iter_unset_bits(..).count(),
+            vob.iter_unset_bits(..).filter(|_| true).count()
+        );
+        if len > 2 {
+            // trigger the edge cases of count_set_bits()
+            let range = 1..len - 1;
+            assert_eq!(
+                vob.iter_set_bits(range.clone()).count(),
+                vob.iter_set_bits(range.clone()).filter(|_| true).count()
+            );
+            assert_eq!(
+                vob.iter_unset_bits(range.clone()).count(),
+                vob.iter_unset_bits(range.clone()).filter(|_| true).count()
+            );
+        }
         vob
     }
 
@@ -2047,4 +2145,37 @@ mod tests {
         v.push(true);
         assert_eq!(v.vec.len(), 1);
     }
+
+    #[test]
+    fn test_count() {
+        let mut rng = rand::rng();
+
+        for test_len in 1..128 {
+            let vob = random_vob(test_len);
+            assert_eq!(
+                vob.iter_storage().count(),
+                vob.iter_storage().filter(|_| true).count()
+            );
+            assert_eq!(vob.iter().count(), vob.iter().filter(|_| true).count());
+            for i in 1..test_len - 1 {
+                let from = rng.random_range(0..i);
+                let to = rng.random_range(from..i);
+                assert_eq!(
+                    vob.iter_set_bits(from..to).count(),
+                    vob.iter_set_bits(from..to).filter(|_| true).count()
+                );
+                assert_eq!(
+                    vob.iter_unset_bits(from..to).count(),
+                    vob.iter_unset_bits(from..to).filter(|_| true).count()
+                );
+            }
+        }
+    }
+
+    #[test]
+    fn test_collect_capacity() {
+        // a test to make sure that iter_set_bits().collect() does not always allocate .len() elements
+        let vec: Vec<usize> = Vob::from_elem(false, 100).iter_set_bits(..).collect();
+        assert_eq!(vec.capacity(), 0);
+    }
 }