Add HashSet::drain_filter method

Fixes #178.

Author: mbrubeck

CHANGELOG.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 ## [Unreleased]
 
+### Added
+- Added a `drain_filter` function to `HashSet`. (#179)
+
 ## [v0.8.0] - 2020-06-18
 
 ### Fixed

src/map.rs

@@ -609,8 +609,10 @@ impl<K, V, S> HashMap<K, V, S> {
     {
         DrainFilter {
             f,
-            iter: unsafe { self.table.iter() },
-            table: &mut self.table,
+            inner: DrainFilterInner {
+                iter: unsafe { self.table.iter() },
+                table: &mut self.table,
+            },
         }
     }
 
@@ -1331,45 +1333,55 @@ where
     F: FnMut(&K, &mut V) -> bool,
 {
     f: F,
-    iter: RawIter<(K, V)>,
-    table: &'a mut RawTable<(K, V)>,
+    inner: DrainFilterInner<'a, K, V>,
 }
 
 impl<'a, K, V, F> Drop for DrainFilter<'a, K, V, F>
 where
     F: FnMut(&K, &mut V) -> bool,
 {
     fn drop(&mut self) {
-        struct DropGuard<'r, 'a, K, V, F>(&'r mut DrainFilter<'a, K, V, F>)
-        where
-            F: FnMut(&K, &mut V) -> bool;
-
-        impl<'r, 'a, K, V, F> Drop for DropGuard<'r, 'a, K, V, F>
-        where
-            F: FnMut(&K, &mut V) -> bool,
-        {
-            fn drop(&mut self) {
-                while let Some(_) = self.0.next() {}
-            }
-        }
         while let Some(item) = self.next() {
-            let guard = DropGuard(self);
+            let guard = ConsumeAllOnDrop(self);
             drop(item);
             mem::forget(guard);
         }
     }
 }
 
+pub(super) struct ConsumeAllOnDrop<'a, T: Iterator>(pub &'a mut T);
+
+impl<T: Iterator> Drop for ConsumeAllOnDrop<'_, T> {
+    fn drop(&mut self) {
+        self.0.for_each(drop)
+    }
+}
+
 impl<K, V, F> Iterator for DrainFilter<'_, K, V, F>
 where
     F: FnMut(&K, &mut V) -> bool,
 {
     type Item = (K, V);
     fn next(&mut self) -> Option<Self::Item> {
+        self.inner.next(&mut self.f)
+    }
+}
+
+/// Portions of `DrainFilter` shared with `set::DrainFilter`
+pub(super) struct DrainFilterInner<'a, K, V> {
+    pub iter: RawIter<(K, V)>,
+    pub table: &'a mut RawTable<(K, V)>,
+}
+
+impl<K, V> DrainFilterInner<'_, K, V> {
+    pub(super) fn next<F>(&mut self, f: &mut F) -> Option<(K, V)>
+    where
+        F: FnMut(&K, &mut V) -> bool,
+    {
         unsafe {
             while let Some(item) = self.iter.next() {
                 let &mut (ref key, ref mut value) = item.as_mut();
-                if !(self.f)(key, value) {
+                if !f(key, value) {
                     return Some(self.table.remove(item));
                 }
             }

src/set.rs

@@ -4,9 +4,10 @@ use core::borrow::Borrow;
 use core::fmt;
 use core::hash::{BuildHasher, Hash};
 use core::iter::{Chain, FromIterator, FusedIterator};
+use core::mem;
 use core::ops::{BitAnd, BitOr, BitXor, Sub};
 
-use super::map::{self, DefaultHashBuilder, HashMap, Keys};
+use super::map::{self, ConsumeAllOnDrop, DefaultHashBuilder, DrainFilterInner, HashMap, Keys};
 
 // Future Optimization (FIXME!)
 // =============================
@@ -285,6 +286,38 @@ impl<T, S> HashSet<T, S> {
         self.map.retain(|k, _| f(k));
     }
 
+    /// Drains elements which are false under the given predicate,
+    /// and returns an iterator over the removed items.
+    ///
+    /// In other words, move all elements `e` such that `f(&e)` returns `false` out
+    /// into another iterator.
+    ///
+    /// When the returned DrainedFilter is dropped, the elements that don't satisfy
+    /// the predicate are dropped from the set.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use hashbrown::HashSet;
+    ///
+    /// let mut set: HashSet<i32> = (0..8).collect();
+    /// let drained = set.drain_filter(|&k| k % 2 == 0);
+    /// assert_eq!(drained.count(), 4);
+    /// assert_eq!(set.len(), 4);
+    /// ```
+    pub fn drain_filter<F>(&mut self, f: F) -> DrainFilter<'_, T, F>
+    where
+        F: FnMut(&T) -> bool,
+    {
+        DrainFilter {
+            f,
+            inner: DrainFilterInner {
+                iter: unsafe { self.map.table.iter() },
+                table: &mut self.map.table,
+            },
+        }
+    }
+
     /// Clears the set, removing all values.
     ///
     /// # Examples
@@ -1185,6 +1218,21 @@ pub struct Drain<'a, K> {
     iter: map::Drain<'a, K, ()>,
 }
 
+/// A draining iterator over entries of a `HashSet` which don't satisfy the predicate `f`.
+///
+/// This `struct` is created by the [`drain_filter`] method on [`HashSet`]. See its
+/// documentation for more.
+///
+/// [`drain_filter`]: struct.HashSet.html#method.drain_filter
+/// [`HashSet`]: struct.HashSet.html
+pub struct DrainFilter<'a, K, F>
+where
+    F: FnMut(&K) -> bool,
+{
+    f: F,
+    inner: DrainFilterInner<'a, K, ()>,
+}
+
 /// A lazy iterator producing elements in the intersection of `HashSet`s.
 ///
 /// This `struct` is created by the [`intersection`] method on [`HashSet`].
@@ -1365,6 +1413,31 @@ impl<K: fmt::Debug> fmt::Debug for Drain<'_, K> {
     }
 }
 
+impl<'a, K, F> Drop for DrainFilter<'a, K, F>
+where
+    F: FnMut(&K) -> bool,
+{
+    fn drop(&mut self) {
+        while let Some(item) = self.next() {
+            let guard = ConsumeAllOnDrop(self);
+            drop(item);
+            mem::forget(guard);
+        }
+    }
+}
+
+impl<K, F> Iterator for DrainFilter<'_, K, F>
+where
+    F: FnMut(&K) -> bool,
+{
+    type Item = K;
+    fn next(&mut self) -> Option<Self::Item> {
+        let f = &mut self.f;
+        let (k, _) = self.inner.next(&mut |k, _| f(k))?;
+        Some(k)
+    }
+}
+
 impl<T, S> Clone for Intersection<'_, T, S> {
     #[cfg_attr(feature = "inline-more", inline)]
     fn clone(&self) -> Self {
@@ -1973,4 +2046,21 @@ mod test_set {
         assert!(set.contains(&4));
         assert!(set.contains(&6));
     }
+
+    #[test]
+    fn test_drain_filter() {
+        {
+            let mut set: HashSet<i32> = (0..8).collect();
+            let drained = set.drain_filter(|&k| k % 2 == 0);
+            let mut out = drained.collect::<Vec<_>>();
+            out.sort_unstable();
+            assert_eq!(vec![1, 3, 5, 7], out);
+            assert_eq!(set.len(), 4);
+        }
+        {
+            let mut set: HashSet<i32> = (0..8).collect();
+            drop(set.drain_filter(|&k| k % 2 == 0));
+            assert_eq!(set.len(), 4, "Removes non-matching items on drop");
+        }
+    }
 }