From 9362c0d0a29c76a00068b200c2aa316ef402a567 Mon Sep 17 00:00:00 2001 From: Christoph Hommelsheim Date: Thu, 10 Sep 2026 14:40:47 +0200 Subject: [PATCH] feat(retain): add `retain` method adds a `retain(&mut self, f: FnMut(u32) -> bool)` method to `RoaringBitmap` and `RoaringTreemap`. It works similar to `Vec::retain`. --- roaring/src/bitmap/container.rs | 6 ++ roaring/src/bitmap/inherent.rs | 107 +++++++++++++++++++++ roaring/src/bitmap/store/bitmap_store.rs | 21 ++++ roaring/src/bitmap/store/interval_store.rs | 22 +++++ roaring/src/bitmap/store/mod.rs | 9 ++ roaring/src/treemap/inherent.rs | 84 ++++++++++++++++ 6 files changed, 249 insertions(+) diff --git a/roaring/src/bitmap/container.rs b/roaring/src/bitmap/container.rs index b11a6446..ce792fcd 100644 --- a/roaring/src/bitmap/container.rs +++ b/roaring/src/bitmap/container.rs @@ -185,6 +185,12 @@ impl Container { }; } + /// Retains only the elements specified by the predicate. + pub fn retain(&mut self, f: impl FnMut(u16) -> bool) { + self.store.retain(f); + self.ensure_correct_store(); + } + pub fn contains(&self, index: u16) -> bool { self.store.contains(index) } diff --git a/roaring/src/bitmap/inherent.rs b/roaring/src/bitmap/inherent.rs index f0b88501..c8d01917 100644 --- a/roaring/src/bitmap/inherent.rs +++ b/roaring/src/bitmap/inherent.rs @@ -856,6 +856,28 @@ impl RoaringBitmap { } } + /// Retains only the elements specified by the predicate. + /// + /// In other word, remove all elements `e` such that `f(e)` returns `false`. + /// + /// # Examples + /// + /// ```rust + /// use roaring::RoaringBitmap; + /// + /// let mut rb = RoaringBitmap::from_iter(0..10); + /// rb.retain(|x| x % 2 == 0); + /// assert_eq!(rb, RoaringBitmap::from_iter([0, 2, 4, 6, 8])); + /// ``` + #[inline] + pub fn retain(&mut self, mut f: impl FnMut(u32) -> bool) { + for container in &mut self.containers { + let key = container.key; + container.retain(|i| f(util::join(key, i))); + } + self.containers.retain(|c| !c.is_empty()); + } + /// Optimizes the container storage for this bitmap. /// Returns true if the container storage was modified, false if not. /// @@ -1191,4 +1213,89 @@ mod tests { bitmap.remove_biggest(4); assert_eq!(bitmap, RoaringBitmap::default()); } + + #[test] + fn retain_array() { + let mut bitmap = RoaringBitmap::from_iter([1, 2, 3, 7, 9, 11]); + bitmap.retain(|x| x % 2 == 1); + assert_eq!(bitmap, RoaringBitmap::from_iter([1, 3, 7, 9, 11])); + + bitmap = RoaringBitmap::from_iter([2, 4, 6, 8]); + bitmap.retain(|x| x % 4 == 0); + assert_eq!(bitmap, RoaringBitmap::from_iter([4, 8])); + + // Retain all + bitmap = RoaringBitmap::from_iter([1, 2, 3]); + bitmap.retain(|_| true); + assert_eq!(bitmap, RoaringBitmap::from_iter([1, 2, 3])); + + // Retain none + bitmap = RoaringBitmap::from_iter([1, 2, 3]); + bitmap.retain(|_| false); + assert!(bitmap.is_empty()); + } + + #[test] + fn retain_bitmap() { + let mut bitmap = RoaringBitmap::new(); + bitmap.insert_range(0..4098); + bitmap.retain(|x| x % 2 == 0); + let expected: RoaringBitmap = (0..4098).step_by(2).collect(); + assert_eq!(bitmap, expected); + + // Retain all + bitmap = RoaringBitmap::new(); + bitmap.insert_range(0..10_000); + bitmap.retain(|_| true); + assert_eq!(bitmap.len(), 10_000); + + // Retain none + bitmap = RoaringBitmap::new(); + bitmap.insert_range(0..10_000); + bitmap.retain(|_| false); + assert!(bitmap.is_empty()); + } + + #[test] + fn retain_multi_container() { + let mut bitmap = RoaringBitmap::from_iter(0u32..65536); + for i in (65536..131072).step_by(2) { + bitmap.insert(i); + } + bitmap.retain(|x| x % 4 == 0); + + let expected: RoaringBitmap = (0..131072).filter(|&x| x % 4 == 0).collect(); + assert_eq!(bitmap, expected); + } + + #[test] + fn retain_empty() { + let mut bitmap = RoaringBitmap::new(); + bitmap.retain(|_| true); + assert!(bitmap.is_empty()); + } + + proptest! { + #[test] + fn retain_preserves_values(mut bitmap in RoaringBitmap::arbitrary()) { + let predicate = |x: u32| !x.is_multiple_of(3); + let expected: RoaringBitmap = bitmap.iter().filter(|&x| predicate(x)).collect(); + + bitmap.retain(predicate); + + prop_assert_eq!(bitmap, expected); + } + + #[test] + fn retain_all_preserves(mut bitmap in RoaringBitmap::arbitrary()) { + bitmap.retain(|_| true); + prop_assert_eq!(bitmap.iter().collect::>(), bitmap.iter().collect::>()); + } + + #[test] + fn retain_none_becomes_empty(mut bitmap in RoaringBitmap::arbitrary()) { + bitmap.retain(|_| false); + prop_assert!(bitmap.is_empty()); + } + } } diff --git a/roaring/src/bitmap/store/bitmap_store.rs b/roaring/src/bitmap/store/bitmap_store.rs index ca32e2b5..c7f6cbd5 100644 --- a/roaring/src/bitmap/store/bitmap_store.rs +++ b/roaring/src/bitmap/store/bitmap_store.rs @@ -405,6 +405,27 @@ impl BitmapStore { } } + /// Retains only the elements specified by the predicate. + pub fn retain(&mut self, mut f: impl FnMut(u16) -> bool) { + let mut new_len = 0u64; + for (key, word) in self.bits.iter_mut().enumerate() { + let mut new_word = 0u64; + let mut remaining = *word; + while remaining != 0 { + let bit = remaining.trailing_zeros(); + let index = (key as u16) * 64 + bit as u16; + if f(index) { + new_word |= 1 << bit; + new_len += 1; + } + // Clear this bit + remaining &= remaining - 1; + } + *word = new_word; + } + self.len = new_len; + } + /// Set N bits that are currently 1 bit from the lower bit to 0. pub fn remove_biggest(&mut self, mut clear_bits: u64) { if self.len() < clear_bits { diff --git a/roaring/src/bitmap/store/interval_store.rs b/roaring/src/bitmap/store/interval_store.rs index 78e010a5..587b9861 100644 --- a/roaring/src/bitmap/store/interval_store.rs +++ b/roaring/src/bitmap/store/interval_store.rs @@ -300,6 +300,28 @@ impl IntervalStore { self.0.drain(remove_to..); } + /// Retains only the elements specified by the predicate. + pub fn retain(&mut self, mut f: impl FnMut(u16) -> bool) { + let mut new_intervals: Vec = Vec::new(); + for interval in self.0.iter() { + // Build retained intervals from this interval + let mut start: Option = None; + for value in *interval { + if f(value) { + if start.is_none() { + start = Some(value); + } + } else if let Some(s) = start.take() { + new_intervals.push(Interval::new_unchecked(s, value - 1)); + } + } + if let Some(s) = start { + new_intervals.push(Interval::new_unchecked(s, interval.end)); + } + } + self.0 = new_intervals; + } + pub fn contains(&self, index: u16) -> bool { self.0.binary_search_by(|iv| cmp_index_interval(index, *iv).reverse()).is_ok() } diff --git a/roaring/src/bitmap/store/mod.rs b/roaring/src/bitmap/store/mod.rs index c96e6045..1110905c 100644 --- a/roaring/src/bitmap/store/mod.rs +++ b/roaring/src/bitmap/store/mod.rs @@ -176,6 +176,15 @@ impl Store { } } + /// Retains only the elements specified by the predicate. + pub fn retain(&mut self, f: impl FnMut(u16) -> bool) { + match self { + Array(vec) => vec.retain(f), + Bitmap(bits) => bits.retain(f), + Run(runs) => runs.retain(f), + } + } + pub fn contains(&self, index: u16) -> bool { match self { Array(vec) => vec.contains(index), diff --git a/roaring/src/treemap/inherent.rs b/roaring/src/treemap/inherent.rs index da8e797d..5eef376f 100644 --- a/roaring/src/treemap/inherent.rs +++ b/roaring/src/treemap/inherent.rs @@ -462,6 +462,26 @@ impl RoaringTreemap { changed } + /// Retains only the elements specified by the predicate. + /// + /// In other word, remove all elements `e` such that `f(e)` returns `false`. + /// + /// # Examples + /// + /// ```rust + /// use roaring::RoaringTreemap; + /// + /// let mut rb = RoaringTreemap::from_iter(0..10); + /// rb.retain(|x| x % 2 == 0); + /// assert_eq!(rb, RoaringTreemap::from_iter([0, 2, 4, 6, 8])); + /// ``` + pub fn retain(&mut self, mut f: impl FnMut(u64) -> bool) { + for (&key, bitmap) in &mut self.map { + bitmap.retain(|i| f(((key as u64) << 32) | i as u64)); + } + self.map.retain(|_, bitmap| !bitmap.is_empty()); + } + /// Returns the minimum value in the set (if the set is non-empty). /// /// # Examples @@ -613,5 +633,69 @@ mod test { prop_assert!(changed); } } + + #[test] + fn retain_preserves_values(mut treemap in RoaringTreemap::arbitrary()) { + let predicate = |x: u64| !x.is_multiple_of(3); + let expected: RoaringTreemap = treemap.iter().filter(|&x| predicate(x)).collect(); + + treemap.retain(predicate); + + prop_assert_eq!(treemap, expected); + } + + #[test] + fn retain_all_preserves(mut treemap in RoaringTreemap::arbitrary()) { + treemap.retain(|_| true); + prop_assert_eq!(treemap.iter().collect::>(), treemap.iter().collect::>()); + } + + #[test] + fn retain_none_becomes_empty(mut treemap in RoaringTreemap::arbitrary()) { + treemap.retain(|_| false); + prop_assert!(treemap.is_empty()); + } + } + + #[test] + fn retain_single_bitmap() { + let mut treemap = RoaringTreemap::from_iter(0..1000u64); + treemap.retain(|x| x % 2 == 0); + let expected: RoaringTreemap = (0..1000).step_by(2).collect(); + assert_eq!(treemap, expected); + } + + #[test] + fn retain_multi_bitmap() { + let mut treemap = RoaringTreemap::new(); + // Fill bitmaps across multiple u32 partitions + for i in 0u64..200_000 { + treemap.insert(i); + } + treemap.retain(|x| x % 7 == 0); + + let expected: RoaringTreemap = (0..200_000).filter(|&x| x % 7 == 0).collect(); + assert_eq!(treemap, expected); + } + + #[test] + fn retain_empty() { + let mut treemap = RoaringTreemap::new(); + treemap.retain(|_| true); + assert!(treemap.is_empty()); + } + + #[test] + fn retain_all() { + let mut treemap = RoaringTreemap::from_iter(0..10_000u64); + treemap.retain(|_| true); + assert_eq!(treemap.len(), 10_000); + } + + #[test] + fn retain_none() { + let mut treemap = RoaringTreemap::from_iter(0..10_000u64); + treemap.retain(|_| false); + assert!(treemap.is_empty()); } }