diff --git a/benches/vob.rs b/benches/vob.rs index 365133b..0b2332e 100644 --- a/benches/vob.rs +++ b/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::())); + v1.extend((0..N).map(|_| rng.random::())); let mut v2 = Vob::with_capacity(N); - v2.extend((0..N).map(|_| rng.gen::())); + v2.extend((0..N).map(|_| rng.random::())); 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::())); + v1.extend((0..N).map(|_| rng.random::())); let mut v2 = Vob::with_capacity(N); - v2.extend((0..N).map(|_| rng.gen::())); + v2.extend((0..N).map(|_| rng.random::())); 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::())); + v1.extend((0..N).map(|_| rng.random::())); let mut v2 = Vob::with_capacity(N); - v2.extend((0..N).map(|_| rng.gen::())); + v2.extend((0..N).map(|_| rng.random::())); c.bench_function("and", |b| { b.iter(|| { @@ -105,7 +105,7 @@ 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))); @@ -114,15 +114,33 @@ fn from_bytes(c: &mut Criterion) { 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::())); - c.bench_function("iter_set_bits", |b| b.iter(|| a.iter_set_bits(..).count())); + a.extend((0..N).map(|_| rng.random::())); + 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::())); + 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::::new_with_storage_type(N); let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED); - a.extend((0..N).map(|_| rng.gen::())); + a.extend((0..N).map(|_| rng.random::())); 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::::new_with_storage_type(N); + let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED); + a.extend((0..N).map(|_| rng.random::())); + c.bench_function("count_set_bits_u8", |b| { b.iter(|| a.iter_set_bits(..).count()) }); } @@ -130,6 +148,13 @@ fn iter_set_bits_u8(c: &mut Criterion) { 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()) }); } @@ -137,6 +162,13 @@ fn iter_all_set_bits(c: &mut Criterion) { 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()) }); } @@ -153,8 +185,12 @@ criterion_group!( and, from_bytes, 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); diff --git a/src/lib.rs b/src/lib.rs index ee72bec..7538451 100644 --- a/src/lib.rs +++ b/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,58 @@ impl Vob { } } + /// Counts the number of set bits. + /// This method assumes the range is processed with process_range() + fn count_set_bits(&self, range: Range) -> usize { + // Early return for empty ranges + if range.is_empty() { + return 0; + } + let start_off = block_offset::(range.start); + + // this -1 arithmetic is safe since we already tested for range.start & range.end equality + let end_off = blocks_required::(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::(); + let end_bit = range.end % bits_per_block::(); + + // 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::() - 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::(); + 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::(); + 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::() - 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 +1135,10 @@ impl Iterator for Iter<'_, T> { fn size_hint(&self) -> (usize, Option) { self.range.size_hint() } + + fn count(self) -> usize { + self.range.count() + } } impl DoubleEndedIterator for Iter<'_, T> { @@ -1148,6 +1205,10 @@ impl Iterator for IterSetBits<'_, T> { fn size_hint(&self) -> (usize, Option) { self.range.size_hint() } + + fn count(self) -> usize { + self.vob.count_set_bits(self.range) + } } impl DoubleEndedIterator for IterSetBits<'_, T> { @@ -1228,6 +1289,12 @@ impl Iterator for IterUnsetBits<'_, T> { fn size_hint(&self) -> (usize, Option) { 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 DoubleEndedIterator for IterUnsetBits<'_, T> { @@ -1300,6 +1367,10 @@ impl Iterator for StorageIter<'_, T> { fn size_hint(&self) -> (usize, Option) { self.iter.size_hint() } + + fn count(self) -> usize { + self.iter.count() + } } #[inline(always)] @@ -1974,6 +2045,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 +2139,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 = Vob::from_elem(false, 100).iter_set_bits(..).collect(); + assert_eq!(vec.capacity(), 0); + } }