Merge #141

141: Add iter size hints r=Kerollmops a=saik0

This is a blocker for #140 as proptest calls `size_hint` when generate / shrinking.

Could only implement ExactSizeIterator on the outer (Bit|Tree)Maps on 64 bit architectures to ensure the `u64` value is not truncated when casting to usize.

Co-authored-by: saik0 <[email protected]>

src/bitmap/container.rs

@@ -237,7 +237,13 @@ impl<'a> Iterator for Iter<'a> {
         self.inner.next().map(|i| util::join(self.key, i))
     }
     fn size_hint(&self) -> (usize, Option<usize>) {
-        panic!("Should never be called (roaring::Iter caches the size_hint itself)")
+        self.inner.size_hint()
+    }
+}
+
+impl<'a> ExactSizeIterator for Iter<'a> {
+    fn len(&self) -> usize {
+        self.inner.len()
     }
 }
 

src/bitmap/iter.rs

@@ -52,6 +52,13 @@ impl Iterator for Iter<'_> {
     }
 }
 
+#[cfg(target_pointer_width = "64")]
+impl ExactSizeIterator for Iter<'_> {
+    fn len(&self) -> usize {
+        self.size_hint as usize
+    }
+}
+
 impl Iterator for IntoIter {
     type Item = u32;
 
@@ -69,6 +76,13 @@ impl Iterator for IntoIter {
     }
 }
 
+#[cfg(target_pointer_width = "64")]
+impl ExactSizeIterator for IntoIter {
+    fn len(&self) -> usize {
+        self.size_hint as usize
+    }
+}
+
 impl RoaringBitmap {
     /// Iterator over each value stored in the RoaringBitmap, guarantees values are ordered by value.
     ///

src/bitmap/store/bitmap_store.rs

@@ -218,11 +218,11 @@ impl BitmapStore {
     }
 
     pub fn iter(&self) -> BitmapIter<&[u64; BITMAP_LENGTH]> {
-        BitmapIter::new(&self.bits)
+        BitmapIter::new(self.len, &self.bits)
     }
 
     pub fn into_iter(self) -> BitmapIter<Box<[u64; BITMAP_LENGTH]>> {
-        BitmapIter::new(self.bits)
+        BitmapIter::new(self.len, self.bits)
     }
 
     pub fn as_array(&self) -> &[u64; BITMAP_LENGTH] {
@@ -261,12 +261,13 @@ impl std::error::Error for Error {}
 pub struct BitmapIter<B: Borrow<[u64; BITMAP_LENGTH]>> {
     key: usize,
     value: u64,
+    len: u64,
     bits: B,
 }
 
 impl<B: Borrow<[u64; BITMAP_LENGTH]>> BitmapIter<B> {
-    fn new(bits: B) -> BitmapIter<B> {
-        BitmapIter { key: 0, value: bits.borrow()[0], bits }
+    fn new(len: u64, bits: B) -> BitmapIter<B> {
+        BitmapIter { key: 0, value: bits.borrow()[0], len, bits }
     }
 }
 
@@ -278,19 +279,27 @@ impl<B: Borrow<[u64; BITMAP_LENGTH]>> Iterator for BitmapIter<B> {
             if self.value == 0 {
                 self.key += 1;
                 if self.key >= BITMAP_LENGTH {
+                    debug_assert!(self.len == 0);
                     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;
+            self.len -= 1;
             return Some((64 * self.key + index) as u16);
         }
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        panic!("Should never be called (roaring::Iter caches the size_hint itself)")
+        (self.len as usize, Some(self.len as usize))
+    }
+}
+
+impl<B: Borrow<[u64; BITMAP_LENGTH]>> ExactSizeIterator for BitmapIter<B> {
+    fn len(&self) -> usize {
+        self.len as usize
     }
 }
 

src/bitmap/store/mod.rs

@@ -430,6 +430,22 @@ impl<'a> Iterator for Iter<'a> {
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        panic!("Should never be called (roaring::Iter caches the size_hint itself)")
+        match self {
+            Iter::Array(inner) => inner.size_hint(),
+            Iter::Vec(inner) => inner.size_hint(),
+            Iter::BitmapBorrowed(inner) => inner.size_hint(),
+            Iter::BitmapOwned(inner) => inner.size_hint(),
+        }
+    }
+}
+
+impl<'a> ExactSizeIterator for Iter<'a> {
+    fn len(&self) -> usize {
+        match self {
+            Iter::Array(inner) => inner.len(),
+            Iter::Vec(inner) => inner.len(),
+            Iter::BitmapBorrowed(inner) => inner.len(),
+            Iter::BitmapOwned(inner) => inner.len(),
+        }
     }
 }

src/treemap/iter.rs

@@ -95,6 +95,13 @@ impl<'a> Iterator for Iter<'a> {
     }
 }
 
+#[cfg(target_pointer_width = "64")]
+impl ExactSizeIterator for Iter<'_> {
+    fn len(&self) -> usize {
+        self.size_hint as usize
+    }
+}
+
 impl Iterator for IntoIter {
     type Item = u64;
 
@@ -112,6 +119,13 @@ impl Iterator for IntoIter {
     }
 }
 
+#[cfg(target_pointer_width = "64")]
+impl ExactSizeIterator for IntoIter {
+    fn len(&self) -> usize {
+        self.size_hint as usize
+    }
+}
+
 impl RoaringTreemap {
     /// Iterator over each value stored in the RoaringTreemap, guarantees values are ordered by
     /// value.