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

(#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

@@ -585,6 +585,57 @@ 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.start >= range.end {
+            return 0;
+        }
+        let start_word = block_offset::<T>(range.start);
+        if start_word >= self.len {
+            return 0;
+        }
+        // this -1 is safe since we already tested for range.start & range.end equality
+        let end_word = blocks_required::<T>(range.end) - 1;
+
+        if start_word == end_word {
+            // Range entirely within one word
+            let word = self.vec[start_word];
+            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
+            return count_ones_usize(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
+                word >> start_bit
+            } else {
+                // We want bits from start_bit to end_bit
+                // After right shift, we need to remove the high bits
+                (word >> start_bit) << (start_bit + bits_per_block::<T>() - end_bit)
+            });
+        }
+
+        // First word: shift out bits before start_bit
+        let start_bit = range.start % bits_per_block::<T>();
+        let mut count = count_ones_usize(self.vec[start_word] >> start_bit);
+
+        // Middle words
+        for word_idx in (start_word + 1)..end_word {
+            count += count_ones_usize(self.vec[word_idx]);
+        }
+
+        // Last word: shift out bits after end_bit
+        let end_bit = range.end % bits_per_block::<T>();
+        if end_bit == 0 {
+            // end_bit = 0 means we want to count the entire end_word
+            count + count_ones_usize(self.vec[end_word])
+        } else {
+            count + count_ones_usize(self.vec[end_word] << (bits_per_block::<T>() - end_bit))
+        }
+    }
+
     /// 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 +1133,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 +1203,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 +1287,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 math is safe since (self.range.end - self.range.start) is the total number of bits.
+        // and self.vob.count_set_bits() always returns less or equal than that.
+        (self.range.end - self.range.start) - self.vob.count_set_bits(self.range)
+    }
 }
 
 impl<T: Debug + PrimInt> DoubleEndedIterator for IterUnsetBits<'_, T> {
@@ -1300,6 +1365,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)]
@@ -1327,6 +1396,16 @@ fn blocks_required<T>(num_bits: usize) -> usize {
     num_bits / bits_per_block::<T>() + usize::from(num_bits % bits_per_block::<T>() != 0)
 }
 
+/// Convenience function that calls T::count_ones() and converts the result to usize
+/// (The conversion is always safe even if T is u128 and usize is u16)
+#[inline(always)]
+fn count_ones_usize<T: PrimInt>(value: T) -> usize {
+    use std::convert::TryFrom;
+    // this could safely use unwrap_unchecked() or ever ’as usize', because the result of
+    // count_ones() will always fit inside an usize. Even if T==u128 and usize == u8.
+    usize::try_from(value.count_ones()).unwrap()
+}
+
 #[macro_export]
 /// Create a `Vob` from a list of boolean values.
 ///
@@ -1974,6 +2053,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 +2147,28 @@ 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);
+            for i in 1..test_len - 1 {
+                let from = rng.random_range(0..i);
+                if from == i {
+                    continue;
+                }
+                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()
+                );
+            }
+        }
+    }
 }