Skip to content

Commit

Permalink
Rollup merge of rust-lang#79026 - mbrubeck:btree_retain, r=m-ou-se
Browse files Browse the repository at this point in the history
Implement BTreeMap::retain and BTreeSet::retain

Adds new methods `BTreeMap::retain` and `BTreeSet::retain`.  These are implemented on top of `drain_filter` (rust-lang#70530).

The API of these methods is identical to `HashMap::retain` and `HashSet::retain`, which were implemented in rust-lang#39560 and stabilized in rust-lang#36648.  The docs and tests are also copied from HashMap/HashSet.

The new methods are unstable, behind the `btree_retain` feature gate, with tracking issue rust-lang#79025.  See also rust-lang/rfcs#1338.
  • Loading branch information
Dylan-DPC authored Nov 13, 2020
2 parents 0180045 + bf6902c commit bb160af
Show file tree
Hide file tree
Showing 4 changed files with 70 additions and 0 deletions.
24 changes: 24 additions & 0 deletions library/alloc/src/collections/btree/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
11 changes: 11 additions & 0 deletions library/alloc/src/collections/btree/map/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::*;

Expand Down
24 changes: 24 additions & 0 deletions library/alloc/src/collections/btree/set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
11 changes: 11 additions & 0 deletions library/alloc/src/collections/btree/set/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down

0 comments on commit bb160af

Please sign in to comment.