extract_if with traits

Author: clarfonthey

src/map.rs

@@ -1,4 +1,4 @@
-use crate::raw::{Allocator, Bucket, Global, RawExtractIf};
+use crate::raw::{Allocator, Bucket, Global};
 use crate::{table, DefaultHashBuilder, Equivalent, HashTable, TryReserveError};
 use ::alloc::borrow::ToOwned;
 use core::borrow::Borrow;
@@ -960,14 +960,10 @@ impl<K, V, S, A: Allocator> HashMap<K, V, S, A> {
     #[cfg_attr(feature = "inline-more", inline)]
     pub fn extract_if<F>(&mut self, f: F) -> ExtractIf<'_, K, V, F, A>
     where
-        F: FnMut(&K, &mut V) -> bool,
+        F: Filter<K, V>,
     {
         ExtractIf {
-            f,
-            inner: RawExtractIf {
-                iter: unsafe { self.table.raw.iter() },
-                table: &mut self.table.raw,
-            },
+            inner: self.table.extract_if(KeyVal(f)),
         }
     }
 
@@ -2590,6 +2586,29 @@ impl<K, V, A: Allocator> Drain<'_, K, V, A> {
     }
 }
 
+/// Adapter between [`map::Filter`](Filter) and [`table::Filter`].
+struct KeyVal<F>(F);
+impl<K, V, F: Filter<K, V>> table::Filter<(K, V)> for KeyVal<F> {
+    #[inline]
+    fn should_extract(&mut self, (ref key, ref mut val): &mut (K, V)) -> bool {
+        self.0.should_extract(key, val)
+    }
+}
+
+/// Filter for [`ExtractIf`].
+///
+/// Accepts `FnMut(&K, &mut V) -> bool`, but can be implemented directly.
+pub trait Filter<K, V> {
+    /// Whether the element should be extracted.
+    fn should_extract(&mut self, key: &K, value: &mut V) -> bool;
+}
+impl<K, V, F: FnMut(&K, &mut V) -> bool> Filter<K, V> for F {
+    #[inline]
+    fn should_extract(&mut self, key: &K, value: &mut V) -> bool {
+        (self)(key, value)
+    }
+}
+
 /// A draining iterator over entries of a `HashMap` which don't satisfy the predicate
 /// `f(&k, &mut v)` in arbitrary order. The iterator element type is `(K, V)`.
 ///
@@ -2623,25 +2642,24 @@ impl<K, V, A: Allocator> Drain<'_, K, V, A> {
 /// ```
 #[must_use = "Iterators are lazy unless consumed"]
 pub struct ExtractIf<'a, K, V, F, A: Allocator = Global> {
-    f: F,
-    inner: RawExtractIf<'a, (K, V), A>,
+    inner: table::ExtractIf<'a, (K, V), KeyVal<F>, A>,
 }
 
 impl<K, V, F, A> Iterator for ExtractIf<'_, K, V, F, A>
 where
-    F: FnMut(&K, &mut V) -> bool,
+    F: Filter<K, V>,
     A: Allocator,
 {
     type Item = (K, V);
 
     #[cfg_attr(feature = "inline-more", inline)]
     fn next(&mut self) -> Option<Self::Item> {
-        self.inner.next(|&mut (ref k, ref mut v)| (self.f)(k, v))
+        self.inner.next()
     }
 
     #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
-        (0, self.inner.iter.size_hint().1)
+        self.inner.size_hint()
     }
 }
 

src/set.rs

@@ -5,8 +5,8 @@ use core::ops::{BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign,
 use core::{fmt, mem};
 use map::make_hash;
 
-use super::map::{self, HashMap, Keys};
-use crate::raw::{Allocator, Global, RawExtractIf};
+use crate::map::{self, HashMap, Keys};
+use crate::raw::{Allocator, Global};
 use crate::DefaultHashBuilder;
 
 // Future Optimization (FIXME!)
@@ -403,14 +403,10 @@ impl<T, S, A: Allocator> HashSet<T, S, A> {
     #[cfg_attr(feature = "inline-more", inline)]
     pub fn extract_if<F>(&mut self, f: F) -> ExtractIf<'_, T, F, A>
     where
-        F: FnMut(&T) -> bool,
+        F: Filter<T>,
     {
         ExtractIf {
-            f,
-            inner: RawExtractIf {
-                iter: unsafe { self.map.table.raw.iter() },
-                table: &mut self.map.table.raw,
-            },
+            inner: self.map.extract_if(Key(f)),
         }
     }
 
@@ -1663,6 +1659,29 @@ pub struct Drain<'a, K, A: Allocator = Global> {
     iter: map::Drain<'a, K, (), A>,
 }
 
+/// Adapter between [`set::Filter`](Filter) and [`map::Filter`].
+struct Key<F>(F);
+impl<T, F: Filter<T>> map::Filter<T, ()> for Key<F> {
+    #[inline]
+    fn should_extract(&mut self, key: &T, (): &mut ()) -> bool {
+        self.0.should_extract(key)
+    }
+}
+
+/// Filter for [`ExtractIf`].
+///
+/// Accepts `FnMut(&K) -> bool`, but can be implemented directly.
+pub trait Filter<K> {
+    /// Whether the element should be extracted.
+    fn should_extract(&mut self, key: &K) -> bool;
+}
+impl<K, F: FnMut(&K) -> bool> Filter<K> for F {
+    #[inline]
+    fn should_extract(&mut self, key: &K) -> bool {
+        (self)(key)
+    }
+}
+
 /// A draining iterator over entries of a `HashSet` which don't satisfy the predicate `f`.
 ///
 /// This `struct` is created by the [`extract_if`] method on [`HashSet`]. See its
@@ -1672,8 +1691,7 @@ pub struct Drain<'a, K, A: Allocator = Global> {
 /// [`HashSet`]: struct.HashSet.html
 #[must_use = "Iterators are lazy unless consumed"]
 pub struct ExtractIf<'a, K, F, A: Allocator = Global> {
-    f: F,
-    inner: RawExtractIf<'a, (K, ()), A>,
+    inner: map::ExtractIf<'a, K, (), Key<F>, A>,
 }
 
 /// A lazy iterator producing elements in the intersection of `HashSet`s.
@@ -1906,20 +1924,18 @@ impl<K: fmt::Debug, A: Allocator> fmt::Debug for Drain<'_, K, A> {
 
 impl<K, F, A: Allocator> Iterator for ExtractIf<'_, K, F, A>
 where
-    F: FnMut(&K) -> bool,
+    F: Filter<K>,
 {
     type Item = K;
 
     #[cfg_attr(feature = "inline-more", inline)]
     fn next(&mut self) -> Option<Self::Item> {
-        self.inner
-            .next(|&mut (ref k, ())| (self.f)(k))
-            .map(|(k, ())| k)
+        self.inner.next().map(|(k, ())| k)
     }
 
     #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
-        (0, self.inner.iter.size_hint().1)
+        self.inner.size_hint()
     }
 }
 

src/table.rs

@@ -1381,7 +1381,7 @@ where
     /// ```
     pub fn extract_if<F>(&mut self, f: F) -> ExtractIf<'_, T, F, A>
     where
-        F: FnMut(&mut T) -> bool,
+        F: Filter<T>,
     {
         ExtractIf {
             f,
@@ -3157,6 +3157,20 @@ impl<T: fmt::Debug, A: Allocator> fmt::Debug for Drain<'_, T, A> {
     }
 }
 
+/// Filter for [`ExtractIf`].
+///
+/// Accepts `FnMut(&mut T) -> bool`, but can be implemented directly.
+pub trait Filter<T> {
+    /// Whether the element should be extracted.
+    fn should_extract(&mut self, elem: &mut T) -> bool;
+}
+impl<T, F: FnMut(&mut T) -> bool> Filter<T> for F {
+    #[inline]
+    fn should_extract(&mut self, elem: &mut T) -> bool {
+        (self)(elem)
+    }
+}
+
 /// A draining iterator over entries of a `HashTable` which don't satisfy the predicate `f`.
 ///
 /// This `struct` is created by [`HashTable::extract_if`]. See its
@@ -3167,15 +3181,12 @@ pub struct ExtractIf<'a, T, F, A: Allocator = Global> {
     inner: RawExtractIf<'a, T, A>,
 }
 
-impl<T, F, A: Allocator> Iterator for ExtractIf<'_, T, F, A>
-where
-    F: FnMut(&mut T) -> bool,
-{
+impl<T, F: Filter<T>, A: Allocator> Iterator for ExtractIf<'_, T, F, A> {
     type Item = T;
 
     #[inline]
     fn next(&mut self) -> Option<Self::Item> {
-        self.inner.next(|val| (self.f)(val))
+        self.inner.next(|val| self.f.should_extract(val))
     }
 
     #[inline]