Skip to content

Commit ba47524

Browse files
eadfeadf
committed
Optimized iter*().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
1 parent 921cfd8 commit ba47524

2 files changed

Lines changed: 183 additions & 11 deletions

File tree

‎benches/vob.rs‎

Lines changed: 65 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -65,9 +65,9 @@ fn split_off(c: &mut Criterion) {
6565
fn xor(c: &mut Criterion) {
6666
let mut v1 = Vob::with_capacity(N);
6767
let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
68-
v1.extend((0..N).map(|_| rng.gen::<bool>()));
68+
v1.extend((0..N).map(|_| rng.random::<bool>()));
6969
let mut v2 = Vob::with_capacity(N);
70-
v2.extend((0..N).map(|_| rng.gen::<bool>()));
70+
v2.extend((0..N).map(|_| rng.random::<bool>()));
7171

7272
c.bench_function("xor", |b| {
7373
b.iter(|| {
@@ -79,9 +79,9 @@ fn xor(c: &mut Criterion) {
7979
fn or(c: &mut Criterion) {
8080
let mut v1 = Vob::with_capacity(N);
8181
let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
82-
v1.extend((0..N).map(|_| rng.gen::<bool>()));
82+
v1.extend((0..N).map(|_| rng.random::<bool>()));
8383
let mut v2 = Vob::with_capacity(N);
84-
v2.extend((0..N).map(|_| rng.gen::<bool>()));
84+
v2.extend((0..N).map(|_| rng.random::<bool>()));
8585

8686
c.bench_function("or", |b| {
8787
b.iter(|| {
@@ -93,9 +93,9 @@ fn or(c: &mut Criterion) {
9393
fn and(c: &mut Criterion) {
9494
let mut v1 = Vob::with_capacity(N);
9595
let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
96-
v1.extend((0..N).map(|_| rng.gen::<bool>()));
96+
v1.extend((0..N).map(|_| rng.random::<bool>()));
9797
let mut v2 = Vob::with_capacity(N);
98-
v2.extend((0..N).map(|_| rng.gen::<bool>()));
98+
v2.extend((0..N).map(|_| rng.random::<bool>()));
9999

100100
c.bench_function("and", |b| {
101101
b.iter(|| {
@@ -105,38 +105,86 @@ fn and(c: &mut Criterion) {
105105
}
106106

107107
fn from_bytes(c: &mut Criterion) {
108-
let mut rng = rand::thread_rng();
108+
let mut rng = rand::rng();
109109
let mut v1 = [0u8; 1024];
110110
rng.fill(&mut v1);
111111
c.bench_function("from_bytes", |b| b.iter(|| Vob::from_bytes(&v1)));
112112
}
113113

114+
fn iter_bits(c: &mut Criterion) {
115+
let mut a = Vob::with_capacity(N);
116+
let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
117+
a.extend((0..N).map(|_| rng.random::<bool>()));
118+
c.bench_function("iter_bits", |b| {
119+
b.iter(|| a.iter().filter(|_| true).count())
120+
});
121+
}
122+
123+
fn count_bits(c: &mut Criterion) {
124+
let mut a = Vob::with_capacity(N);
125+
let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
126+
a.extend((0..N).map(|_| rng.random::<bool>()));
127+
c.bench_function("count_bits", |b| b.iter(|| a.iter().count()));
128+
}
129+
114130
fn iter_set_bits(c: &mut Criterion) {
115131
let mut a = Vob::with_capacity(N);
116132
let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
117-
a.extend((0..N).map(|_| rng.gen::<bool>()));
118-
c.bench_function("iter_set_bits", |b| b.iter(|| a.iter_set_bits(..).count()));
133+
a.extend((0..N).map(|_| rng.random::<bool>()));
134+
c.bench_function("iter_set_bits", |b| {
135+
b.iter(|| a.iter_set_bits(..).filter(|_| true).count())
136+
});
137+
}
138+
139+
fn count_set_bits(c: &mut Criterion) {
140+
let mut a = Vob::with_capacity(N);
141+
let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
142+
a.extend((0..N).map(|_| rng.random::<bool>()));
143+
c.bench_function("count_set_bits", |b| b.iter(|| a.iter_set_bits(..).count()));
119144
}
120145

121146
fn iter_set_bits_u8(c: &mut Criterion) {
122147
let mut a = Vob::<u8>::new_with_storage_type(N);
123148
let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
124-
a.extend((0..N).map(|_| rng.gen::<bool>()));
149+
a.extend((0..N).map(|_| rng.random::<bool>()));
125150
c.bench_function("iter_set_bits_u8", |b| {
151+
b.iter(|| a.iter_set_bits(..).filter(|_| true).count())
152+
});
153+
}
154+
155+
fn count_set_bits_u8(c: &mut Criterion) {
156+
let mut a = Vob::<u8>::new_with_storage_type(N);
157+
let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
158+
a.extend((0..N).map(|_| rng.random::<bool>()));
159+
c.bench_function("count_set_bits_u8", |b| {
126160
b.iter(|| a.iter_set_bits(..).count())
127161
});
128162
}
129163

130164
fn iter_all_set_bits(c: &mut Criterion) {
131165
let a = Vob::from_elem(true, N);
132166
c.bench_function("iter_all_set_bits", |b| {
167+
b.iter(|| a.iter_set_bits(..).filter(|_| true).count())
168+
});
169+
}
170+
171+
fn count_all_set_bits(c: &mut Criterion) {
172+
let a = Vob::from_elem(true, N);
173+
c.bench_function("count_all_set_bits", |b| {
133174
b.iter(|| a.iter_set_bits(..).count())
134175
});
135176
}
136177

137178
fn iter_all_unset_bits(c: &mut Criterion) {
138179
let a = Vob::from_elem(true, N);
139180
c.bench_function("iter_all_unset_bits", |b| {
181+
b.iter(|| a.iter_unset_bits(..).filter(|_| true).count())
182+
});
183+
}
184+
185+
fn count_all_unset_bits(c: &mut Criterion) {
186+
let a = Vob::from_elem(true, N);
187+
c.bench_function("count_all_unset_bits", |b| {
140188
b.iter(|| a.iter_unset_bits(..).count())
141189
});
142190
}
@@ -152,9 +200,15 @@ criterion_group!(
152200
or,
153201
and,
154202
from_bytes,
203+
iter_bits,
204+
count_bits,
155205
iter_set_bits,
206+
count_set_bits,
156207
iter_set_bits_u8,
208+
count_set_bits_u8,
157209
iter_all_set_bits,
158-
iter_all_unset_bits
210+
count_all_set_bits,
211+
iter_all_unset_bits,
212+
count_all_unset_bits
159213
);
160214
criterion_main!(benches);

‎src/lib.rs‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -585,6 +585,56 @@ impl<T: Debug + PrimInt> Vob<T> {
585585
}
586586
}
587587

588+
/// Counts the number of set bits.
589+
/// This method assumes the range is processed with process_range()
590+
fn _count_set_bits(&self, range: Range<usize>) -> usize {
591+
// Early return for empty ranges
592+
if range.start >= range.end {
593+
return 0;
594+
}
595+
let start_word = block_offset::<T>(range.start);
596+
if start_word >= self.len {
597+
return 0;
598+
}
599+
// this -1 is safe since we already tested for range.start & range.end equality
600+
let end_word = blocks_required::<T>(range.end) - 1;
601+
602+
if start_word == end_word {
603+
// Range entirely within one word
604+
let word = self.vec[start_word];
605+
let start_bit = range.start % bits_per_block::<T>();
606+
let end_bit = range.end % bits_per_block::<T>();
607+
608+
// Remove bits before start_bit and bits after end_bit
609+
return count_ones_usize(if end_bit == 0 {
610+
// end_bit = 0 means we want everything from start_bit to end of word
611+
// After the right shift above, we already have what we want
612+
word >> start_bit
613+
} else {
614+
// We want bits from start_bit to end_bit
615+
// After right shift, we need to remove the high bits
616+
(word >> start_bit) << (start_bit + bits_per_block::<T>() - end_bit)
617+
});
618+
}
619+
620+
// First word: shift out bits before start_bit
621+
let start_bit = range.start % bits_per_block::<T>();
622+
let mut count = count_ones_usize(self.vec[start_word] >> start_bit);
623+
624+
// Middle words (unchanged)
625+
for word_idx in (start_word + 1)..end_word {
626+
count += count_ones_usize(self.vec[word_idx]);
627+
}
628+
629+
// Last word: shift out bits after end_bit
630+
let end_bit = range.end % bits_per_block::<T>();
631+
if end_bit == 0 {
632+
count + count_ones_usize(self.vec[end_word])
633+
} else {
634+
count + count_ones_usize(self.vec[end_word] << (bits_per_block::<T>() - end_bit))
635+
}
636+
}
637+
588638
/// Returns an iterator which efficiently produces the index of each unset bit in the specified
589639
/// range. Assuming appropriate support from your CPU, this is much more efficient than
590640
/// checking each bit individually.
@@ -1082,6 +1132,10 @@ impl<T: Debug + PrimInt> Iterator for Iter<'_, T> {
10821132
fn size_hint(&self) -> (usize, Option<usize>) {
10831133
self.range.size_hint()
10841134
}
1135+
1136+
fn count(self) -> usize {
1137+
self.range.count()
1138+
}
10851139
}
10861140

10871141
impl<T: Debug + PrimInt> DoubleEndedIterator for Iter<'_, T> {
@@ -1148,6 +1202,10 @@ impl<T: Debug + PrimInt> Iterator for IterSetBits<'_, T> {
11481202
fn size_hint(&self) -> (usize, Option<usize>) {
11491203
self.range.size_hint()
11501204
}
1205+
1206+
fn count(self) -> usize {
1207+
self.vob._count_set_bits(self.range)
1208+
}
11511209
}
11521210

11531211
impl<T: Debug + PrimInt> DoubleEndedIterator for IterSetBits<'_, T> {
@@ -1228,6 +1286,10 @@ impl<T: Debug + PrimInt> Iterator for IterUnsetBits<'_, T> {
12281286
fn size_hint(&self) -> (usize, Option<usize>) {
12291287
self.range.size_hint()
12301288
}
1289+
1290+
fn count(self) -> usize {
1291+
(self.range.end - self.range.start) - self.vob._count_set_bits(self.range)
1292+
}
12311293
}
12321294

12331295
impl<T: Debug + PrimInt> DoubleEndedIterator for IterUnsetBits<'_, T> {
@@ -1300,6 +1362,10 @@ impl<T: Debug + PrimInt> Iterator for StorageIter<'_, T> {
13001362
fn size_hint(&self) -> (usize, Option<usize>) {
13011363
self.iter.size_hint()
13021364
}
1365+
1366+
fn count(self) -> usize {
1367+
self.iter.count()
1368+
}
13031369
}
13041370

13051371
#[inline(always)]
@@ -1330,6 +1396,14 @@ fn blocks_required<T>(num_bits: usize) -> usize {
13301396
}
13311397
}
13321398

1399+
/// Convenience function that calls T::count_ones() and converts the result to usize
1400+
/// (The conversion is always safe even if T is u128 and usize is u16)
1401+
#[inline(always)]
1402+
fn count_ones_usize<T: PrimInt>(value: T) -> usize {
1403+
use std::convert::TryFrom;
1404+
usize::try_from(value.count_ones()).unwrap()
1405+
}
1406+
13331407
#[macro_export]
13341408
/// Create a `Vob` from a list of boolean values.
13351409
///
@@ -1977,6 +2051,26 @@ mod tests {
19772051
for _ in 0..len {
19782052
vob.push(rng.random());
19792053
}
2054+
assert_eq!(
2055+
vob.iter_set_bits(..).count(),
2056+
vob.iter_set_bits(..).filter(|_| true).count()
2057+
);
2058+
assert_eq!(
2059+
vob.iter_unset_bits(..).count(),
2060+
vob.iter_unset_bits(..).filter(|_| true).count()
2061+
);
2062+
if len > 2 {
2063+
// trigger the edge cases of _count_set_bits()
2064+
let range = 1..len - 1;
2065+
assert_eq!(
2066+
vob.iter_set_bits(range.clone()).count(),
2067+
vob.iter_set_bits(range.clone()).filter(|_| true).count()
2068+
);
2069+
assert_eq!(
2070+
vob.iter_unset_bits(range.clone()).count(),
2071+
vob.iter_unset_bits(range.clone()).filter(|_| true).count()
2072+
);
2073+
}
19802074
vob
19812075
}
19822076

@@ -2050,4 +2144,28 @@ mod tests {
20502144
v.push(true);
20512145
assert_eq!(v.vec.len(), 1);
20522146
}
2147+
2148+
#[test]
2149+
fn test_count() {
2150+
let mut rng = rand::rng();
2151+
2152+
for test_len in 1..128 {
2153+
let vob = random_vob(test_len);
2154+
for i in 1..test_len - 1 {
2155+
let from = rng.random_range(0..i);
2156+
if from == i {
2157+
continue;
2158+
}
2159+
let to = rng.random_range(from..i);
2160+
assert_eq!(
2161+
vob.iter_set_bits(from..to).count(),
2162+
vob.iter_set_bits(from..to).filter(|_| true).count()
2163+
);
2164+
assert_eq!(
2165+
vob.iter_unset_bits(from..to).count(),
2166+
vob.iter_unset_bits(from..to).filter(|_| true).count()
2167+
);
2168+
}
2169+
}
2170+
}
20532171
}

0 commit comments

Comments
 (0)