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 BTreeMap { } } + /// 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 = (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(&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 = (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 BTreeSet { 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 = 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(&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 = 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();