add dedicated fold and rfold impls

add test for `fold` also add test for treemap
RoaringBitmap · May 29, 2024 · eb26e0d · eb26e0d
1 parent 3fa2e73
commit eb26e0d
Show file tree

Hide file tree

Showing 4 changed files with 152 additions and 0 deletions.
diff --git a/src/bitmap/iter.rs b/src/bitmap/iter.rs
@@ -46,13 +46,30 @@ impl Iterator for Iter<'_> {
             (usize::MAX, None)
         }
     }
+
+    #[inline]
+    fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        self.inner.fold(init, f)
+    }
 }
 
 impl DoubleEndedIterator for Iter<'_> {
     fn next_back(&mut self) -> Option<Self::Item> {
         self.size_hint = self.size_hint.saturating_sub(1);
         self.inner.next_back()
     }
+
+    #[inline]
+    fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
+    where
+        Fold: FnMut(Acc, Self::Item) -> Acc,
+    {
+        self.inner.rfold(init, fold)
+    }
 }
 
 #[cfg(target_pointer_width = "64")]
@@ -77,13 +94,30 @@ impl Iterator for IntoIter {
             (usize::MAX, None)
         }
     }
+
+    #[inline]
+    fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        self.inner.fold(init, f)
+    }
 }
 
 impl DoubleEndedIterator for IntoIter {
     fn next_back(&mut self) -> Option<Self::Item> {
         self.size_hint = self.size_hint.saturating_sub(1);
         self.inner.next_back()
     }
+
+    #[inline]
+    fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
+    where
+        Fold: FnMut(Acc, Self::Item) -> Acc,
+    {
+        self.inner.rfold(init, fold)
+    }
 }
 
 #[cfg(target_pointer_width = "64")]

diff --git a/src/treemap/iter.rs b/src/treemap/iter.rs
@@ -16,12 +16,40 @@ impl<'a> Iterator for To64Iter<'a> {
     fn next(&mut self) -> Option<u64> {
         self.inner.next().map(|n| util::join(self.hi, n))
     }
+
+    #[inline]
+    fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        #[inline]
+        fn f_with_hi<B>(mut f: impl FnMut(B, u64) -> B, hi: u32) -> impl FnMut(B, u32) -> B {
+            move |b, lo| f(b, ((hi as u64) << 32) + (lo as u64))
+        }
+
+        self.inner.fold(init, f_with_hi(f, self.hi))
+    }
 }
 
 impl DoubleEndedIterator for To64Iter<'_> {
     fn next_back(&mut self) -> Option<Self::Item> {
         self.inner.next_back().map(|n| util::join(self.hi, n))
     }
+
+    #[inline]
+    fn rfold<B, F>(self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        #[inline]
+        fn f_with_hi<B>(mut f: impl FnMut(B, u64) -> B, hi: u32) -> impl FnMut(B, u32) -> B {
+            move |b, lo| f(b, ((hi as u64) << 32) + (lo as u64))
+        }
+
+        self.inner.rfold(init, f_with_hi(f, self.hi))
+    }
 }
 
 fn to64iter<'a>(t: (&'a u32, &'a RoaringBitmap)) -> To64Iter<'a> {
@@ -38,12 +66,40 @@ impl Iterator for To64IntoIter {
     fn next(&mut self) -> Option<u64> {
         self.inner.next().map(|n| util::join(self.hi, n))
     }
+
+    #[inline]
+    fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        #[inline]
+        fn f_with_hi<B>(mut f: impl FnMut(B, u64) -> B, hi: u32) -> impl FnMut(B, u32) -> B {
+            move |b, lo| f(b, ((hi as u64) << 32) + (lo as u64))
+        }
+
+        self.inner.fold(init, f_with_hi(f, self.hi))
+    }
 }
 
 impl DoubleEndedIterator for To64IntoIter {
     fn next_back(&mut self) -> Option<Self::Item> {
         self.inner.next_back().map(|n| util::join(self.hi, n))
     }
+
+    #[inline]
+    fn rfold<B, F>(self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        #[inline]
+        fn f_with_hi<B>(mut f: impl FnMut(B, u64) -> B, hi: u32) -> impl FnMut(B, u32) -> B {
+            move |b, lo| f(b, ((hi as u64) << 32) + (lo as u64))
+        }
+
+        self.inner.rfold(init, f_with_hi(f, self.hi))
+    }
 }
 
 fn to64intoiter(t: (u32, RoaringBitmap)) -> To64IntoIter {
@@ -104,13 +160,30 @@ impl<'a> Iterator for Iter<'a> {
             (usize::MAX, None)
         }
     }
+
+    #[inline]
+    fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        self.inner.fold(init, f)
+    }
 }
 
 impl DoubleEndedIterator for Iter<'_> {
     fn next_back(&mut self) -> Option<Self::Item> {
         self.size_hint = self.size_hint.saturating_sub(1);
         self.inner.next_back()
     }
+
+    #[inline]
+    fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
+    where
+        Fold: FnMut(Acc, Self::Item) -> Acc,
+    {
+        self.inner.rfold(init, fold)
+    }
 }
 
 #[cfg(target_pointer_width = "64")]
@@ -135,13 +208,30 @@ impl Iterator for IntoIter {
             (usize::MAX, None)
         }
     }
+
+    #[inline]
+    fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        self.inner.fold(init, f)
+    }
 }
 
 impl DoubleEndedIterator for IntoIter {
     fn next_back(&mut self) -> Option<Self::Item> {
         self.size_hint = self.size_hint.saturating_sub(1);
         self.inner.next_back()
     }
+
+    #[inline]
+    fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
+    where
+        Fold: FnMut(Acc, Self::Item) -> Acc,
+    {
+        self.inner.rfold(init, fold)
+    }
 }
 
 #[cfg(target_pointer_width = "64")]

diff --git a/tests/iter.rs b/tests/iter.rs
@@ -67,6 +67,20 @@ proptest! {
     }
 }
 
+proptest! {
+    #[test]
+    fn fold(values in btree_set(any::<u32>(), ..=10_000)) {
+        let bitmap = RoaringBitmap::from_sorted_iter(values.iter().cloned()).unwrap();
+        let mut val_iter = values.into_iter();
+        // `Iterator::all` uses currently unimplementable `try_fold`, we test `fold`
+        #[allow(clippy::unnecessary_fold)]
+        let r = bitmap.into_iter().fold(true, |b, i| {
+            b && i == val_iter.next().unwrap()
+        });
+        assert!(r)
+    }
+}
+
 #[test]
 fn rev_array() {
     let values = 0..100;

diff --git a/tests/treemap_iter.rs b/tests/treemap_iter.rs
@@ -77,6 +77,20 @@ proptest! {
     }
 }
 
+proptest! {
+    #[test]
+    fn fold(values in btree_set(any::<u64>(), ..=10_000)) {
+        let bitmap = RoaringTreemap::from_sorted_iter(values.iter().cloned()).unwrap();
+        let mut val_iter = values.into_iter();
+        // `Iterator::all` uses currently unimplementable `try_fold`, we test `fold`
+        #[allow(clippy::unnecessary_fold)]
+        let r = bitmap.into_iter().fold(true, |b, i| {
+            b && i == val_iter.next().unwrap()
+        });
+        assert!(r)
+    }
+}
+
 #[test]
 fn rev() {
     let values = (1..3)