feat(retain): add retain method

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)
     }

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::<Vec<_>>(), bitmap.iter().collect::<Vec<_>>());
+        }
+
+        #[test]
+        fn retain_none_becomes_empty(mut bitmap in RoaringBitmap::arbitrary()) {
+            bitmap.retain(|_| false);
+            prop_assert!(bitmap.is_empty());
+        }
+    }
 }

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 {

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<Interval> = Vec::new();
+        for interval in self.0.iter() {
+            // Build retained intervals from this interval
+            let mut start: Option<u16> = 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()
     }

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),

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::<Vec<_>>(), treemap.iter().collect::<Vec<_>>());
+        }
+
+        #[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());
     }
 }