diff --git a/library/std/src/collections/hash/map.rs b/library/std/src/collections/hash/map.rs index fc0175332a769..1a3a493fbb8f6 100644 --- a/library/std/src/collections/hash/map.rs +++ b/library/std/src/collections/hash/map.rs @@ -497,6 +497,50 @@ impl HashMap { Drain { base: self.base.drain() } } + /// Creates an iterator which uses a closure to determine if an element should be removed. + /// + /// If the closure returns true, the element is removed from the map and yielded. + /// If the closure returns false, or panics, the element remains in the map and will not be + /// yielded. + /// + /// Note that `drain_filter` lets you mutate every value in the filter closure, regardless of + /// whether you choose to keep or remove it. + /// + /// If the iterator is only partially consumed or not consumed at all, each of the remaining + /// elements will still be subjected to the closure and removed and dropped if it returns true. + /// + /// It is unspecified how many more elements will be subjected to the closure + /// if a panic occurs in the closure, or a panic occurs while dropping an element, + /// or if the `DrainFilter` value is leaked. + /// + /// # Examples + /// + /// Splitting a map into even and odd keys, reusing the original map: + /// + /// ``` + /// #![feature(hash_drain_filter)] + /// use std::collections::HashMap; + /// + /// let mut map: HashMap = (0..8).map(|x| (x, x)).collect(); + /// let drained: HashMap = map.drain_filter(|k, _v| k % 2 == 0).collect(); + /// + /// let mut evens = drained.keys().copied().collect::>(); + /// let mut odds = map.keys().copied().collect::>(); + /// evens.sort(); + /// odds.sort(); + /// + /// assert_eq!(evens, vec![0, 2, 4, 6]); + /// assert_eq!(odds, vec![1, 3, 5, 7]); + /// ``` + #[inline] + #[unstable(feature = "hash_drain_filter", issue = "59618")] + pub fn drain_filter(&mut self, pred: F) -> DrainFilter<'_, K, V, F> + where + F: FnMut(&K, &mut V) -> bool, + { + DrainFilter { base: self.base.drain_filter(pred) } + } + /// Clears the map, removing all key-value pairs. Keeps the allocated memory /// for reuse. /// @@ -1190,6 +1234,19 @@ impl<'a, K, V> Drain<'a, K, V> { } } +/// A draining, filtering iterator over the entries of a `HashMap`. +/// +/// This `struct` is created by the [`drain_filter`] method on [`HashMap`]. +/// +/// [`drain_filter`]: HashMap::drain_filter +#[unstable(feature = "hash_drain_filter", issue = "59618")] +pub struct DrainFilter<'a, K, V, F> +where + F: FnMut(&K, &mut V) -> bool, +{ + base: base::DrainFilter<'a, K, V, F>, +} + /// A mutable iterator over the values of a `HashMap`. /// /// This `struct` is created by the [`values_mut`] method on [`HashMap`]. See its @@ -1990,6 +2047,36 @@ where } } +#[unstable(feature = "hash_drain_filter", issue = "59618")] +impl Iterator for DrainFilter<'_, K, V, F> +where + F: FnMut(&K, &mut V) -> bool, +{ + type Item = (K, V); + + #[inline] + fn next(&mut self) -> Option<(K, V)> { + self.base.next() + } + #[inline] + fn size_hint(&self) -> (usize, Option) { + self.base.size_hint() + } +} + +#[unstable(feature = "hash_drain_filter", issue = "59618")] +impl FusedIterator for DrainFilter<'_, K, V, F> where F: FnMut(&K, &mut V) -> bool {} + +#[unstable(feature = "hash_drain_filter", issue = "59618")] +impl<'a, K, V, F> fmt::Debug for DrainFilter<'a, K, V, F> +where + F: FnMut(&K, &mut V) -> bool, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.pad("DrainFilter { .. }") + } +} + impl<'a, K, V> Entry<'a, K, V> { #[stable(feature = "rust1", since = "1.0.0")] /// Ensures a value is in the entry by inserting the default if empty, and returns diff --git a/library/std/src/collections/hash/set.rs b/library/std/src/collections/hash/set.rs index f9bc2e933dcf9..72f4798b65d66 100644 --- a/library/std/src/collections/hash/set.rs +++ b/library/std/src/collections/hash/set.rs @@ -247,6 +247,47 @@ impl HashSet { Drain { base: self.base.drain() } } + /// Creates an iterator which uses a closure to determine if a value should be removed. + /// + /// If the closure returns true, then the value is removed and yielded. + /// If the closure returns false, the value will remain in the list and will not be yielded + /// by the iterator. + /// + /// If the iterator is only partially consumed or not consumed at all, each of the remaining + /// values will still be subjected to the closure and removed and dropped if it returns true. + /// + /// It is unspecified how many more values will be subjected to the closure + /// if a panic occurs in the closure, or if a panic occurs while dropping a value, or if the + /// `DrainFilter` itself is leaked. + /// + /// # Examples + /// + /// Splitting a set into even and odd values, reusing the original set: + /// + /// ``` + /// #![feature(hash_drain_filter)] + /// use std::collections::HashSet; + /// + /// let mut set: HashSet = (0..8).collect(); + /// let drained: HashSet = set.drain_filter(|v| v % 2 == 0).collect(); + /// + /// let mut evens = drained.into_iter().collect::>(); + /// let mut odds = set.into_iter().collect::>(); + /// evens.sort(); + /// odds.sort(); + /// + /// assert_eq!(evens, vec![0, 2, 4, 6]); + /// assert_eq!(odds, vec![1, 3, 5, 7]); + /// ``` + #[inline] + #[unstable(feature = "hash_drain_filter", issue = "59618")] + pub fn drain_filter(&mut self, pred: F) -> DrainFilter<'_, T, F> + where + F: FnMut(&T) -> bool, + { + DrainFilter { base: self.base.drain_filter(pred) } + } + /// Clears the set, removing all values. /// /// # Examples @@ -1159,6 +1200,19 @@ pub struct Drain<'a, K: 'a> { base: base::Drain<'a, K>, } +/// A draining, filtering iterator over the items of a `HashSet`. +/// +/// This `struct` is created by the [`drain_filter`] method on [`HashSet`]. +/// +/// [`drain_filter`]: HashSet::drain_filter +#[unstable(feature = "hash_drain_filter", issue = "59618")] +pub struct DrainFilter<'a, K, F> +where + F: FnMut(&K) -> bool, +{ + base: base::DrainFilter<'a, K, F>, +} + /// A lazy iterator producing elements in the intersection of `HashSet`s. /// /// This `struct` is created by the [`intersection`] method on [`HashSet`]. @@ -1348,6 +1402,36 @@ impl fmt::Debug for Drain<'_, K> { } } +#[unstable(feature = "hash_drain_filter", issue = "59618")] +impl Iterator for DrainFilter<'_, K, F> +where + F: FnMut(&K) -> bool, +{ + type Item = K; + + #[inline] + fn next(&mut self) -> Option { + self.base.next() + } + #[inline] + fn size_hint(&self) -> (usize, Option) { + self.base.size_hint() + } +} + +#[unstable(feature = "hash_drain_filter", issue = "59618")] +impl FusedIterator for DrainFilter<'_, K, F> where F: FnMut(&K) -> bool {} + +#[unstable(feature = "hash_drain_filter", issue = "59618")] +impl<'a, K, F> fmt::Debug for DrainFilter<'a, K, F> +where + F: FnMut(&K) -> bool, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.pad("DrainFilter { .. }") + } +} + #[stable(feature = "rust1", since = "1.0.0")] impl Clone for Intersection<'_, T, S> { #[inline]