From bf6902ca61ae6e91a3ab95bed88426926f245d02 Mon Sep 17 00:00:00 2001
From: Matt Brubeck <mbrubeck@limpet.net>
Date: Fri, 13 Nov 2020 10:23:50 -0800
Subject: [PATCH] Add BTreeMap::retain and BTreeSet::retain

---
 library/alloc/src/collections/btree/map.rs    | 24 +++++++++++++++++++
 .../alloc/src/collections/btree/map/tests.rs  | 11 +++++++++
 library/alloc/src/collections/btree/set.rs    | 24 +++++++++++++++++++
 .../alloc/src/collections/btree/set/tests.rs  | 11 +++++++++
 4 files changed, 70 insertions(+)

diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs
index 49122f53d33ad..7151d3763f01f 100644
--- a/library/alloc/src/collections/btree/map.rs
+++ b/library/alloc/src/collections/btree/map.rs
@@ -863,6 +863,30 @@ impl<K: Ord, V> BTreeMap<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 dd3ebcccf76a5..11dbb584abdac 100644
--- 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 684019f8f5f5e..1a807100653bc 100644
--- a/library/alloc/src/collections/btree/set.rs
+++ b/library/alloc/src/collections/btree/set.rs
@@ -798,6 +798,30 @@ impl<T: Ord> BTreeSet<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 52cde8299e418..ef40a048a382e 100644
--- a/library/alloc/src/collections/btree/set/tests.rs
+++ b/library/alloc/src/collections/btree/set/tests.rs
@@ -324,6 +324,17 @@ fn test_is_subset() {
     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();