Add BTreeMap::retain and BTreeSet::retain

diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs
index 49122f53d33ad0a7b2f74be1f9558245d70604b4..7151d3763f01fcb3c0dd7f8572fcf24aa2e49881 100644 (file)
--- a/library/alloc/src/collections/btree/map.rs
+++ b/library/alloc/src/collections/btree/map.rs
@@ -863,6 +863,30 @@ pub fn remove_entry<Q: ?Sized>(&mut self, key: &Q) -> Option<(K, V)>
         }
     }
 
+    /// Retains only the elements specified by the predicate.
+    ///
+    /// In other words, remove all pairs `(k, v)` such that `f(&k, &mut v)` returns `false`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(btree_retain)]
+    /// use std::collections::BTreeMap;
+    ///
+    /// let mut map: BTreeMap<i32, i32> = (0..8).map(|x| (x, x*10)).collect();
+    /// // Keep only the elements with even-numbered keys.
+    /// map.retain(|&k, _| k % 2 == 0);
+    /// assert!(map.into_iter().eq(vec![(0, 0), (2, 20), (4, 40), (6, 60)]));
+    /// ```
+    #[inline]
+    #[unstable(feature = "btree_retain", issue = "79025")]
+    pub fn retain<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&K, &mut V) -> bool,
+    {
+        self.drain_filter(|k, v| !f(k, v));
+    }
+
     /// Moves all elements from `other` into `Self`, leaving `other` empty.
     ///
     /// # Examples

diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs
index dd3ebcccf76a5c131914642d7883763cd21873e3..11dbb584abdac6de774d6c6b5b99b2bc7bf6bd3e 100644 (file)
--- a/library/alloc/src/collections/btree/map/tests.rs
+++ b/library/alloc/src/collections/btree/map/tests.rs
@@ -808,6 +808,17 @@ fn test_range_mut() {
     map.check();
 }
 
+#[test]
+fn test_retain() {
+    let mut map: BTreeMap<i32, i32> = (0..100).map(|x| (x, x * 10)).collect();
+
+    map.retain(|&k, _| k % 2 == 0);
+    assert_eq!(map.len(), 50);
+    assert_eq!(map[&2], 20);
+    assert_eq!(map[&4], 40);
+    assert_eq!(map[&6], 60);
+}
+
 mod test_drain_filter {
     use super::*;
 

diff --git a/library/alloc/src/collections/btree/set.rs b/library/alloc/src/collections/btree/set.rs
index 684019f8f5f5e02c6afa472a5a960686a1d46e59..1a807100653bc1a9270d679dd7107d9d824616de 100644 (file)
--- a/library/alloc/src/collections/btree/set.rs
+++ b/library/alloc/src/collections/btree/set.rs
@@ -798,6 +798,30 @@ pub fn take<Q: ?Sized>(&mut self, value: &Q) -> Option<T>
         Recover::take(&mut self.map, value)
     }
 
+    /// Retains only the elements specified by the predicate.
+    ///
+    /// In other words, remove all elements `e` such that `f(&e)` returns `false`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(btree_retain)]
+    /// use std::collections::BTreeSet;
+    ///
+    /// let xs = [1, 2, 3, 4, 5, 6];
+    /// let mut set: BTreeSet<i32> = xs.iter().cloned().collect();
+    /// // Keep only the even numbers.
+    /// set.retain(|&k| k % 2 == 0);
+    /// assert!(set.iter().eq([2, 4, 6].iter()));
+    /// ```
+    #[unstable(feature = "btree_retain", issue = "79025")]
+    pub fn retain<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&T) -> bool,
+    {
+        self.drain_filter(|v| !f(v));
+    }
+
     /// Moves all elements from `other` into `Self`, leaving `other` empty.
     ///
     /// # Examples

diff --git a/library/alloc/src/collections/btree/set/tests.rs b/library/alloc/src/collections/btree/set/tests.rs
index 52cde8299e41862fc0cf821600dfbf6c4238b4c7..ef40a048a382ecccf8911610c8e9e64e13dc9f18 100644 (file)
--- a/library/alloc/src/collections/btree/set/tests.rs
+++ b/library/alloc/src/collections/btree/set/tests.rs
@@ -324,6 +324,17 @@ fn is_subset(a: &[i32], b: &[i32]) -> bool {
     assert_eq!(is_subset(&[99, 100], &large), false);
 }
 
+#[test]
+fn test_retain() {
+    let xs = [1, 2, 3, 4, 5, 6];
+    let mut set: BTreeSet<i32> = xs.iter().cloned().collect();
+    set.retain(|&k| k % 2 == 0);
+    assert_eq!(set.len(), 3);
+    assert!(set.contains(&2));
+    assert!(set.contains(&4));
+    assert!(set.contains(&6));
+}
+
 #[test]
 fn test_drain_filter() {
     let mut x: BTreeSet<_> = [1].iter().copied().collect();