Skip to content

Commit 560faed

Browse files
eadfeadf
committed
Optimized .count() for the iterators (softdevteam#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 560faed

2 files changed

Lines changed: 204 additions & 11 deletions

File tree

‎benches/vob.rs‎

Lines changed: 67 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,88 @@ 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+
// this benchmark can later be removed, since it is only benchmarking std::Rage.count() as a
124+
// comparison with the old way of doing it. (iter_bits() above)
125+
fn count_bits(c: &mut Criterion) {
126+
let mut a = Vob::with_capacity(N);
127+
let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
128+
a.extend((0..N).map(|_| rng.random::<bool>()));
129+
c.bench_function("count_bits", |b| b.iter(|| a.iter().count()));
130+
}
131+
114132
fn iter_set_bits(c: &mut Criterion) {
115133
let mut a = Vob::with_capacity(N);
116134
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()));
135+
a.extend((0..N).map(|_| rng.random::<bool>()));
136+
c.bench_function("iter_set_bits", |b| {
137+
b.iter(|| a.iter_set_bits(..).filter(|_| true).count())
138+
});
139+
}
140+
141+
fn count_set_bits(c: &mut Criterion) {
142+
let mut a = Vob::with_capacity(N);
143+
let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
144+
a.extend((0..N).map(|_| rng.random::<bool>()));
145+
c.bench_function("count_set_bits", |b| b.iter(|| a.iter_set_bits(..).count()));
119146
}
120147

121148
fn iter_set_bits_u8(c: &mut Criterion) {
122149
let mut a = Vob::<u8>::new_with_storage_type(N);
123150
let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
124-
a.extend((0..N).map(|_| rng.gen::<bool>()));
151+
a.extend((0..N).map(|_| rng.random::<bool>()));
125152
c.bench_function("iter_set_bits_u8", |b| {
153+
b.iter(|| a.iter_set_bits(..).filter(|_| true).count())
154+
});
155+
}
156+
157+
fn count_set_bits_u8(c: &mut Criterion) {
158+
let mut a = Vob::<u8>::new_with_storage_type(N);
159+
let mut rng = Pcg64Mcg::seed_from_u64(RNG_SEED);
160+
a.extend((0..N).map(|_| rng.random::<bool>()));
161+
c.bench_function("count_set_bits_u8", |b| {
126162
b.iter(|| a.iter_set_bits(..).count())
127163
});
128164
}
129165

130166
fn iter_all_set_bits(c: &mut Criterion) {
131167
let a = Vob::from_elem(true, N);
132168
c.bench_function("iter_all_set_bits", |b| {
169+
b.iter(|| a.iter_set_bits(..).filter(|_| true).count())
170+
});
171+
}
172+
173+
fn count_all_set_bits(c: &mut Criterion) {
174+
let a = Vob::from_elem(true, N);
175+
c.bench_function("count_all_set_bits", |b| {
133176
b.iter(|| a.iter_set_bits(..).count())
134177
});
135178
}
136179

137180
fn iter_all_unset_bits(c: &mut Criterion) {
138181
let a = Vob::from_elem(true, N);
139182
c.bench_function("iter_all_unset_bits", |b| {
183+
b.iter(|| a.iter_unset_bits(..).filter(|_| true).count())
184+
});
185+
}
186+
187+
fn count_all_unset_bits(c: &mut Criterion) {
188+
let a = Vob::from_elem(true, N);
189+
c.bench_function("count_all_unset_bits", |b| {
140190
b.iter(|| a.iter_unset_bits(..).count())
141191
});
142192
}
@@ -152,9 +202,15 @@ criterion_group!(
152202
or,
153203
and,
154204
from_bytes,
205+
iter_bits,
206+
count_bits,
155207
iter_set_bits,
208+
count_set_bits,
156209
iter_set_bits_u8,
210+
count_set_bits_u8,
157211
iter_all_set_bits,
158-
iter_all_unset_bits
212+
count_all_set_bits,
213+
iter_all_unset_bits,
214+
count_all_unset_bits
159215
);
160216
criterion_main!(benches);

‎src/lib.rs‎

Lines changed: 137 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,70 @@ 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+
debug_assert!(
598+
start_off < self.vec.len(),
599+
"start_off {} >= self.vec.len() {}",
600+
start_off,
601+
self.vec.len()
602+
);
603+
604+
// this -1 arithmetic is safe since we already tested for range.start & range.end equality
605+
let end_off = blocks_required::<T>(range.end) - 1;
606+
debug_assert!(
607+
end_off < self.vec.len(),
608+
"end_off {} >= self.vec.len() {}",
609+
end_off,
610+
self.vec.len()
611+
);
612+
613+
if start_off == end_off {
614+
// Range entirely within one word
615+
let b = self.vec[start_off];
616+
let start_bit = range.start % bits_per_block::<T>();
617+
let end_bit = range.end % bits_per_block::<T>();
618+
619+
// Remove bits before start_bit and bits after end_bit
620+
let count = if end_bit == 0 {
621+
// end_bit = 0 means we want everything from start_bit to end of word
622+
// After the right shift, we have what we want
623+
b >> start_bit
624+
} else {
625+
// We want bits from start_bit to end_bit
626+
// After right shift, we need to remove the high bits
627+
(b >> start_bit) << (start_bit + bits_per_block::<T>() - end_bit)
628+
}
629+
.count_ones();
630+
return usize::try_from(count).unwrap();
631+
}
632+
633+
// First word: shift out bits before start_bit
634+
let start_bit = range.start % bits_per_block::<T>();
635+
let mut count = usize::try_from((self.vec[start_off] >> start_bit).count_ones()).unwrap();
636+
637+
// Middle words
638+
for word_idx in (start_off + 1)..end_off {
639+
count += usize::try_from(self.vec[word_idx].count_ones()).unwrap();
640+
}
641+
642+
// Last word: shift out bits after end_bit
643+
let end_bit = range.end % bits_per_block::<T>();
644+
let count_ones = if end_bit == 0 {
645+
// end_bit = 0 means we want to count the entire end_off word
646+
self.vec[end_off].count_ones()
647+
} else {
648+
(self.vec[end_off] << (bits_per_block::<T>() - end_bit)).count_ones()
649+
};
650+
count + usize::try_from(count_ones).unwrap()
651+
}
652+
588653
/// Returns an iterator which efficiently produces the index of each unset bit in the specified
589654
/// range. Assuming appropriate support from your CPU, this is much more efficient than
590655
/// checking each bit individually.
@@ -1082,6 +1147,10 @@ impl<T: Debug + PrimInt> Iterator for Iter<'_, T> {
10821147
fn size_hint(&self) -> (usize, Option<usize>) {
10831148
self.range.size_hint()
10841149
}
1150+
1151+
fn count(self) -> usize {
1152+
self.range.count()
1153+
}
10851154
}
10861155

10871156
impl<T: Debug + PrimInt> DoubleEndedIterator for Iter<'_, T> {
@@ -1148,6 +1217,10 @@ impl<T: Debug + PrimInt> Iterator for IterSetBits<'_, T> {
11481217
fn size_hint(&self) -> (usize, Option<usize>) {
11491218
self.range.size_hint()
11501219
}
1220+
1221+
fn count(self) -> usize {
1222+
self.vob.count_set_bits(self.range)
1223+
}
11511224
}
11521225

11531226
impl<T: Debug + PrimInt> DoubleEndedIterator for IterSetBits<'_, T> {
@@ -1228,6 +1301,12 @@ impl<T: Debug + PrimInt> Iterator for IterUnsetBits<'_, T> {
12281301
fn size_hint(&self) -> (usize, Option<usize>) {
12291302
self.range.size_hint()
12301303
}
1304+
1305+
fn count(self) -> usize {
1306+
// This arithmetic is safe because (self.range.end - self.range.start) is the total number of bits,
1307+
// and self.vob.count_set_bits() always returns a value less than or equal to that.
1308+
(self.range.end - self.range.start) - self.vob.count_set_bits(self.range)
1309+
}
12311310
}
12321311

12331312
impl<T: Debug + PrimInt> DoubleEndedIterator for IterUnsetBits<'_, T> {
@@ -1300,6 +1379,10 @@ impl<T: Debug + PrimInt> Iterator for StorageIter<'_, T> {
13001379
fn size_hint(&self) -> (usize, Option<usize>) {
13011380
self.iter.size_hint()
13021381
}
1382+
1383+
fn count(self) -> usize {
1384+
self.iter.count()
1385+
}
13031386
}
13041387

13051388
#[inline(always)]
@@ -1974,6 +2057,27 @@ mod tests {
19742057
for _ in 0..len {
19752058
vob.push(rng.random());
19762059
}
2060+
// these tests can later be dialed down, as they noticeable slow down every random vob test.
2061+
assert_eq!(
2062+
vob.iter_set_bits(..).count(),
2063+
vob.iter_set_bits(..).filter(|_| true).count()
2064+
);
2065+
assert_eq!(
2066+
vob.iter_unset_bits(..).count(),
2067+
vob.iter_unset_bits(..).filter(|_| true).count()
2068+
);
2069+
if len > 2 {
2070+
// trigger the edge cases of count_set_bits()
2071+
let range = 1..len - 1;
2072+
assert_eq!(
2073+
vob.iter_set_bits(range.clone()).count(),
2074+
vob.iter_set_bits(range.clone()).filter(|_| true).count()
2075+
);
2076+
assert_eq!(
2077+
vob.iter_unset_bits(range.clone()).count(),
2078+
vob.iter_unset_bits(range.clone()).filter(|_| true).count()
2079+
);
2080+
}
19772081
vob
19782082
}
19792083

@@ -2047,4 +2151,37 @@ mod tests {
20472151
v.push(true);
20482152
assert_eq!(v.vec.len(), 1);
20492153
}
2154+
2155+
#[test]
2156+
fn test_count() {
2157+
let mut rng = rand::rng();
2158+
2159+
for test_len in 1..128 {
2160+
let vob = random_vob(test_len);
2161+
assert_eq!(
2162+
vob.iter_storage().count(),
2163+
vob.iter_storage().filter(|_| true).count()
2164+
);
2165+
assert_eq!(vob.iter().count(), vob.iter().filter(|_| true).count());
2166+
for i in 1..test_len - 1 {
2167+
let from = rng.random_range(0..i);
2168+
let to = rng.random_range(from..i);
2169+
assert_eq!(
2170+
vob.iter_set_bits(from..to).count(),
2171+
vob.iter_set_bits(from..to).filter(|_| true).count()
2172+
);
2173+
assert_eq!(
2174+
vob.iter_unset_bits(from..to).count(),
2175+
vob.iter_unset_bits(from..to).filter(|_| true).count()
2176+
);
2177+
}
2178+
}
2179+
}
2180+
2181+
#[test]
2182+
fn test_collect_capacity() {
2183+
// a test to make sure that iter_set_bits().collect() does not always allocate .len() elements
2184+
let vec: Vec<usize> = Vob::from_elem(false, 100).iter_set_bits(..).collect();
2185+
assert_eq!(vec.capacity(), 0);
2186+
}
20502187
}

0 commit comments

Comments
 (0)