Skip to content

Commit 1e6912a

Browse files
eadfeadf
committed
Optimized .count() for the iterators (#87)
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 c9c5577 commit 1e6912a

2 files changed

Lines changed: 174 additions & 11 deletions

File tree

‎benches/vob.rs‎

Lines changed: 49 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,7 +105,7 @@ 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)));
@@ -114,29 +114,61 @@ fn from_bytes(c: &mut Criterion) {
114114
fn iter_set_bits(c: &mut Criterion) {
115115
let mut a = Vob::with_capacity(N);
116116
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()));
117+
a.extend((0..N).map(|_| rng.random::<bool>()));
118+
c.bench_function("iter_set_bits", |b| {
119+
b.iter(|| a.iter_set_bits(..).filter(|_| true).count())
120+
});
121+
}
122+
123+
fn count_set_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_set_bits", |b| b.iter(|| a.iter_set_bits(..).count()));
119128
}
120129

121130
fn iter_set_bits_u8(c: &mut Criterion) {
122131
let mut a = Vob::<u8>::new_with_storage_type(N);
123132
let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
124-
a.extend((0..N).map(|_| rng.gen::<bool>()));
133+
a.extend((0..N).map(|_| rng.random::<bool>()));
125134
c.bench_function("iter_set_bits_u8", |b| {
135+
b.iter(|| a.iter_set_bits(..).filter(|_| true).count())
136+
});
137+
}
138+
139+
fn count_set_bits_u8(c: &mut Criterion) {
140+
let mut a = Vob::<u8>::new_with_storage_type(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_u8", |b| {
126144
b.iter(|| a.iter_set_bits(..).count())
127145
});
128146
}
129147

130148
fn iter_all_set_bits(c: &mut Criterion) {
131149
let a = Vob::from_elem(true, N);
132150
c.bench_function("iter_all_set_bits", |b| {
151+
b.iter(|| a.iter_set_bits(..).filter(|_| true).count())
152+
});
153+
}
154+
155+
fn count_all_set_bits(c: &mut Criterion) {
156+
let a = Vob::from_elem(true, N);
157+
c.bench_function("count_all_set_bits", |b| {
133158
b.iter(|| a.iter_set_bits(..).count())
134159
});
135160
}
136161

137162
fn iter_all_unset_bits(c: &mut Criterion) {
138163
let a = Vob::from_elem(true, N);
139164
c.bench_function("iter_all_unset_bits", |b| {
165+
b.iter(|| a.iter_unset_bits(..).filter(|_| true).count())
166+
});
167+
}
168+
169+
fn count_all_unset_bits(c: &mut Criterion) {
170+
let a = Vob::from_elem(true, N);
171+
c.bench_function("count_all_unset_bits", |b| {
140172
b.iter(|| a.iter_unset_bits(..).count())
141173
});
142174
}
@@ -152,9 +184,15 @@ criterion_group!(
152184
or,
153185
and,
154186
from_bytes,
187+
iter_bits,
188+
count_bits,
155189
iter_set_bits,
190+
count_set_bits,
156191
iter_set_bits_u8,
192+
count_set_bits_u8,
157193
iter_all_set_bits,
158-
iter_all_unset_bits
194+
count_all_set_bits,
195+
iter_all_unset_bits,
196+
count_all_unset_bits
159197
);
160198
criterion_main!(benches);

‎src/lib.rs‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
66
use std::{
77
cmp::{min, PartialEq},
8+
convert::TryFrom,
89
fmt::{self, Debug},
910
hash::{Hash, Hasher},
1011
iter::{DoubleEndedIterator, FromIterator, FusedIterator},
@@ -585,6 +586,58 @@ impl<T: Debug + PrimInt> Vob<T> {
585586
}
586587
}
587588

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

10871144
impl<T: Debug + PrimInt> DoubleEndedIterator for Iter<'_, T> {
@@ -1148,6 +1205,10 @@ impl<T: Debug + PrimInt> Iterator for IterSetBits<'_, T> {
11481205
fn size_hint(&self) -> (usize, Option<usize>) {
11491206
self.range.size_hint()
11501207
}
1208+
1209+
fn count(self) -> usize {
1210+
self.vob.count_set_bits(self.range)
1211+
}
11511212
}
11521213

11531214
impl<T: Debug + PrimInt> DoubleEndedIterator for IterSetBits<'_, T> {
@@ -1228,6 +1289,12 @@ impl<T: Debug + PrimInt> Iterator for IterUnsetBits<'_, T> {
12281289
fn size_hint(&self) -> (usize, Option<usize>) {
12291290
self.range.size_hint()
12301291
}
1292+
1293+
fn count(self) -> usize {
1294+
// This arithmetic is safe because (self.range.end - self.range.start) is the total number of bits,
1295+
// and self.vob.count_set_bits() always returns a value less than or equal to that.
1296+
(self.range.end - self.range.start) - self.vob.count_set_bits(self.range)
1297+
}
12311298
}
12321299

12331300
impl<T: Debug + PrimInt> DoubleEndedIterator for IterUnsetBits<'_, T> {
@@ -1300,6 +1367,10 @@ impl<T: Debug + PrimInt> Iterator for StorageIter<'_, T> {
13001367
fn size_hint(&self) -> (usize, Option<usize>) {
13011368
self.iter.size_hint()
13021369
}
1370+
1371+
fn count(self) -> usize {
1372+
self.iter.count()
1373+
}
13031374
}
13041375

13051376
#[inline(always)]
@@ -1974,6 +2045,27 @@ mod tests {
19742045
for _ in 0..len {
19752046
vob.push(rng.random());
19762047
}
2048+
// these tests can later be dialed down, as they noticeable slow down every random vob test.
2049+
assert_eq!(
2050+
vob.iter_set_bits(..).count(),
2051+
vob.iter_set_bits(..).filter(|_| true).count()
2052+
);
2053+
assert_eq!(
2054+
vob.iter_unset_bits(..).count(),
2055+
vob.iter_unset_bits(..).filter(|_| true).count()
2056+
);
2057+
if len > 2 {
2058+
// trigger the edge cases of count_set_bits()
2059+
let range = 1..len - 1;
2060+
assert_eq!(
2061+
vob.iter_set_bits(range.clone()).count(),
2062+
vob.iter_set_bits(range.clone()).filter(|_| true).count()
2063+
);
2064+
assert_eq!(
2065+
vob.iter_unset_bits(range.clone()).count(),
2066+
vob.iter_unset_bits(range.clone()).filter(|_| true).count()
2067+
);
2068+
}
19772069
vob
19782070
}
19792071

@@ -2047,4 +2139,37 @@ mod tests {
20472139
v.push(true);
20482140
assert_eq!(v.vec.len(), 1);
20492141
}
2142+
2143+
#[test]
2144+
fn test_count() {
2145+
let mut rng = rand::rng();
2146+
2147+
for test_len in 1..128 {
2148+
let vob = random_vob(test_len);
2149+
assert_eq!(
2150+
vob.iter_storage().count(),
2151+
vob.iter_storage().filter(|_| true).count()
2152+
);
2153+
assert_eq!(vob.iter().count(), vob.iter().filter(|_| true).count());
2154+
for i in 1..test_len - 1 {
2155+
let from = rng.random_range(0..i);
2156+
let to = rng.random_range(from..i);
2157+
assert_eq!(
2158+
vob.iter_set_bits(from..to).count(),
2159+
vob.iter_set_bits(from..to).filter(|_| true).count()
2160+
);
2161+
assert_eq!(
2162+
vob.iter_unset_bits(from..to).count(),
2163+
vob.iter_unset_bits(from..to).filter(|_| true).count()
2164+
);
2165+
}
2166+
}
2167+
}
2168+
2169+
#[test]
2170+
fn test_collect_capacity() {
2171+
// a test to make sure that iter_set_bits().collect() does not always allocate .len() elements
2172+
let vec: Vec<usize> = Vob::from_elem(false, 100).iter_set_bits(..).collect();
2173+
assert_eq!(vec.capacity(), 0);
2174+
}
20502175
}

0 commit comments

Comments
 (0)