Speed up bitmap store iterator

RoaringBitmap · Oct 16, 2024 · 385eff7 · 385eff7
1 parent c7c3cd5
commit 385eff7
Showing 1 changed file with 19 additions and 16 deletions.
diff --git a/roaring/src/bitmap/store/bitmap_store.rs b/roaring/src/bitmap/store/bitmap_store.rs
@@ -1,5 +1,4 @@
 use core::borrow::Borrow;
-use core::cmp::Ordering;
 use core::fmt::{Display, Formatter};
 use core::ops::{BitAndAssign, BitOrAssign, BitXorAssign, RangeInclusive, SubAssign};
 
@@ -427,24 +426,28 @@ impl<B: Borrow<[u64; BITMAP_LENGTH]>> Iterator for BitmapIter<B> {
     type Item = u16;
 
     fn next(&mut self) -> Option<u16> {
-        loop {
-            if self.value == 0 {
-                self.key += 1;
-                let cmp = self.key.cmp(&self.key_back);
-                // Match arms can be reordered, this ordering is perf sensitive
-                self.value = if cmp == Ordering::Less {
-                    unsafe { *self.bits.borrow().get_unchecked(self.key) }
-                } else if cmp == Ordering::Equal {
-                    self.value_back
-                } else {
+        if self.value == 0 {
+            'get_val: {
+                if self.key >= self.key_back {
                     return None;
-                };
-                continue;
+                }
+                for key in self.key + 1..self.key_back {
+                    self.value = unsafe { *self.bits.borrow().get_unchecked(key) };
+                    if self.value != 0 {
+                        self.key = key;
+                        break 'get_val;
+                    }
+                }
+                self.key = self.key_back;
+                self.value = self.value_back;
+                if self.value == 0 {
+                    return None;
+                }
             }
-            let index = self.value.trailing_zeros() as usize;
-            self.value &= self.value - 1;
-            return Some((64 * self.key + index) as u16);
         }
+        let index = self.value.trailing_zeros() as usize;
+        self.value &= self.value - 1;
+        Some((64 * self.key + index) as u16)
     }
 }