Speed up bitmap iteration

benchmarks/benches/lib.rs

@@ -245,17 +245,64 @@ fn insert_range_bitmap(c: &mut Criterion) {
 }
 
 fn iter(c: &mut Criterion) {
-    c.bench_function("iter", |b| {
+    c.bench_function("iter bitmap 1..10_000", |b| {
         let bitmap: RoaringBitmap = (1..10_000).collect();
+        b.iter(|| {
+            bitmap.iter().for_each(|i| {
+                black_box(i);
+            });
+        });
+    });
+
+    c.bench_function("iter bitmap sparse", |b| {
+        let bitmap: RoaringBitmap = (0..1 << 16).step_by(61).collect();
+        b.iter(|| {
+            bitmap.iter().for_each(|i| {
+                black_box(i);
+            });
+        });
+    });
 
+    c.bench_function("iter bitmap dense", |b| {
+        let bitmap: RoaringBitmap = (0..1 << 16).step_by(2).collect();
         b.iter(|| {
-            let mut sum: u32 = 0;
+            bitmap.iter().for_each(|i| {
+                black_box(i);
+            });
+        });
+    });
 
-            for (_, element) in bitmap.iter().enumerate() {
-                sum += element;
-            }
+    c.bench_function("iter bitmap minimal", |b| {
+        let bitmap: RoaringBitmap = (0..4096).collect();
+        b.iter(|| {
+            bitmap.iter().for_each(|i| {
+                black_box(i);
+            });
+        });
+    });
+
+    c.bench_function("iter bitmap full", |b| {
+        let bitmap: RoaringBitmap = (0..1 << 16).collect();
+        b.iter(|| {
+            bitmap.iter().for_each(|i| {
+                black_box(i);
+            });
+        });
+    });
+
+    c.bench_function("iter parsed", |b| {
+        let files = self::datasets_paths::WIKILEAKS_NOQUOTES_SRT;
+        let parsed_numbers = parse_dir_files(files).unwrap();
 
-            assert_eq!(sum, 49_995_000);
+        let bitmaps: Vec<_> = parsed_numbers
+            .into_iter()
+            .map(|(_, r)| r.map(|iter| RoaringBitmap::from_sorted_iter(iter).unwrap()).unwrap())
+            .collect();
+
+        b.iter(|| {
+            bitmaps.iter().flat_map(|bitmap| bitmap.iter()).for_each(|i| {
+                black_box(i);
+            });
         });
     });
 }

benchmarks/build.rs

@@ -49,6 +49,7 @@ fn main() -> anyhow::Result<()> {
         "// This file is generated by the build script.\n// Do not modify by hand, use the build.rs file.\n"
     )?;
 
+    #[allow(clippy::single_element_loop)]
     for dataset in &[DATASET_WIKILEAKS_NOQUOTES_SRT] {
         let out_path = out_dir.join(dataset);
         let url = format!("{}/{}.zip", BASE_URL, dataset);

src/bitmap/store.rs

@@ -24,7 +24,7 @@ pub enum Iter<'a> {
 
 pub struct BitmapIter<B: Borrow<[u64; BITMAP_LENGTH]>> {
     key: usize,
-    bit: usize,
+    value: u64,
     bits: B,
 }
 
@@ -699,15 +699,7 @@ impl Clone for Store {
 
 impl<B: Borrow<[u64; BITMAP_LENGTH]>> BitmapIter<B> {
     fn new(bits: B) -> BitmapIter<B> {
-        BitmapIter { key: 0, bit: 0, bits }
-    }
-
-    fn move_next(&mut self) {
-        self.bit += 1;
-        if self.bit == 64 {
-            self.bit = 0;
-            self.key += 1;
-        }
+        BitmapIter { key: 0, value: bits.borrow()[0], bits }
     }
 }
 
@@ -716,17 +708,17 @@ impl<B: Borrow<[u64; BITMAP_LENGTH]>> Iterator for BitmapIter<B> {
 
     fn next(&mut self) -> Option<u16> {
         loop {
-            if self.key == BITMAP_LENGTH {
-                return None;
-            } else if (unsafe { self.bits.borrow().get_unchecked(self.key) } & (1u64 << self.bit))
-                != 0
-            {
-                let result = Some((self.key * 64 + self.bit) as u16);
-                self.move_next();
-                return result;
-            } else {
-                self.move_next();
+            if self.value == 0 {
+                self.key += 1;
+                if self.key >= BITMAP_LENGTH {
+                    return None;
+                }
+                self.value = unsafe { *self.bits.borrow().get_unchecked(self.key) };
+                continue;
             }
+            let index = self.value.trailing_zeros() as usize;
+            self.value &= self.value - 1;
+            return Some((64 * self.key + index) as u16);
         }
     }
 

tests/iter.rs

@@ -1,8 +1,12 @@
 extern crate roaring;
-use roaring::RoaringBitmap;
 
+use std::collections::BTreeSet;
 use std::iter::FromIterator;
 
+use quickcheck_macros::quickcheck;
+
+use roaring::RoaringBitmap;
+
 #[test]
 fn array() {
     let original = (0..2000).collect::<RoaringBitmap>();
@@ -48,3 +52,10 @@ fn bitmaps() {
     assert_eq!(clone, original);
     assert_eq!(clone2, original);
 }
+
+#[quickcheck]
+fn qc_iter(values: BTreeSet<u32>) {
+    let bitmap = RoaringBitmap::from_sorted_iter(values.iter().cloned()).unwrap();
+    // Iterator::eq != PartialEq::eq - cannot use assert_eq macro
+    assert!(values.into_iter().eq(bitmap.into_iter()));
+}