From e0c7f13b1c130d54757d06ea05add323e86d339d Mon Sep 17 00:00:00 2001 From: julianknodt Date: Sun, 15 Dec 2019 12:52:52 -0800 Subject: [PATCH 1/4] Implemented drain_filter Drain filter is as described in the doc comments. The implementation is based almost entirely on retain's implementation, as per `amanieu`'s suggestion. I messed around with the lifetimes, as I'm not entirely familiar with the unsafe `iter` on the raw table, but since we're now using a lazy iterator, the predicate must be valid for as long as the borrow on the table. I also annotated the function with a `#[must_use]`, otherwise the drain would have no effect. Please let me know if there are any other additions before this change can be added. Thanks! --- src/map.rs | 42 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/src/map.rs b/src/map.rs index fcd0e53d5a..4a2db964e7 100644 --- a/src/map.rs +++ b/src/map.rs @@ -956,6 +956,40 @@ where } } } + /// Retains only the elements specified by the predicate, and returns an iterator over the + /// removed items. + /// + /// In other words, move all pairs `(k, v)` such that `f(&k,&mut v)` returns `false` out + /// into another iterator. + /// + /// # Examples + /// + /// ``` + /// use hashbrown::HashMap; + /// + /// let mut map: HashMap = (0..8).map(|x|(x, x*10)).collect(); + /// let drained = map.drain_filter(|&k, _| k % 2 == 0); + /// assert_eq!(drained.count(), 4); + /// assert_eq!(map.len(), 4); + /// ``` + pub fn drain_filter<'a, F>(&'a mut self, mut f: F) -> impl Iterator + '_ + where + F: 'a + FnMut(&K, &mut V) -> bool, + { + // Here we only use `iter` as a temporary, preventing use-after-free + unsafe { + self.table.iter().filter_map(move |item| { + let &mut (ref key, ref mut value) = item.as_mut(); + if !f(key, value) { + // Erase the element from the table first since drop might panic. + self.table.erase_no_drop(&item); + Some(item.read()) + } else { + None + } + }) + } + } } impl HashMap { @@ -3488,6 +3522,14 @@ mod test_map { assert_eq!(map[&6], 60); } + #[test] + fn test_drain_filter() { + let mut map: HashMap = (0..8).map(|x|(x, x*10)).collect(); + let drained = map.drain_filter(|&k, _| k % 2 == 0); + assert_eq!(drained.count(), 4); + assert_eq!(map.len(), 4); + } + #[test] #[cfg_attr(miri, ignore)] // FIXME: no OOM signalling (https://github.com/rust-lang/miri/issues/613) fn test_try_reserve() { From 4dd50ab00bbbe4e94d98b585599c926fb755e171 Mon Sep 17 00:00:00 2001 From: julianknodt Date: Sun, 15 Dec 2019 16:11:05 -0800 Subject: [PATCH 2/4] Updated to fix clippy Fixed clippy lint (added tab and spaced out `=(K,V)` => `= (K, V)`). Also made the test more explicit. Edit: I've finally installed clippy rather than checking Travis. --- src/map.rs | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/src/map.rs b/src/map.rs index 4a2db964e7..cb8e76ea9e 100644 --- a/src/map.rs +++ b/src/map.rs @@ -956,8 +956,8 @@ where } } } - /// Retains only the elements specified by the predicate, and returns an iterator over the - /// removed items. + /// Drains elements which are false under the given predicate, + /// and returns an iterator over the removed items. /// /// In other words, move all pairs `(k, v)` such that `f(&k,&mut v)` returns `false` out /// into another iterator. @@ -972,7 +972,7 @@ where /// assert_eq!(drained.count(), 4); /// assert_eq!(map.len(), 4); /// ``` - pub fn drain_filter<'a, F>(&'a mut self, mut f: F) -> impl Iterator + '_ + pub fn drain_filter<'a, F>(&'a mut self, mut f: F) -> impl Iterator + '_ where F: 'a + FnMut(&K, &mut V) -> bool, { @@ -980,12 +980,11 @@ where unsafe { self.table.iter().filter_map(move |item| { let &mut (ref key, ref mut value) = item.as_mut(); - if !f(key, value) { - // Erase the element from the table first since drop might panic. + if f(key, value) { + None + } else { self.table.erase_no_drop(&item); Some(item.read()) - } else { - None } }) } @@ -3524,10 +3523,12 @@ mod test_map { #[test] fn test_drain_filter() { - let mut map: HashMap = (0..8).map(|x|(x, x*10)).collect(); - let drained = map.drain_filter(|&k, _| k % 2 == 0); - assert_eq!(drained.count(), 4); - assert_eq!(map.len(), 4); + let mut map: HashMap = (0..8).map(|x| (x, x * 10)).collect(); + let drained = map.drain_filter(|&k, _| k % 2 == 0); + let mut out = drained.collect::>(); + out.sort_unstable(); + assert_eq!(vec![(1, 10), (3, 30), (5, 50), (7, 70)], out); + assert_eq!(map.len(), 4); } #[test] From 1d5fac25f79e18946eea4b65f0d55baec5129d88 Mon Sep 17 00:00:00 2001 From: julianknodt Date: Sun, 15 Dec 2019 20:41:34 -0800 Subject: [PATCH 3/4] Updated with DrainFilter struct The previous implementation did not match vec's drain filter's drop semantics. As per `amanieu`'s suggestion, added a DrainFilter which implements drop so that it removes the items which don't satisfy the predicate. --- src/map.rs | 93 +++++++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 74 insertions(+), 19 deletions(-) diff --git a/src/map.rs b/src/map.rs index cb8e76ea9e..f3e8c86e1b 100644 --- a/src/map.rs +++ b/src/map.rs @@ -962,6 +962,9 @@ where /// In other words, move all pairs `(k, v)` such that `f(&k,&mut v)` returns `false` out /// into another iterator. /// + /// When the returned DrainedFilter is dropped, the elements that don't satisfy + /// the predicate are dropped from the table. + /// /// # Examples /// /// ``` @@ -972,21 +975,14 @@ where /// assert_eq!(drained.count(), 4); /// assert_eq!(map.len(), 4); /// ``` - pub fn drain_filter<'a, F>(&'a mut self, mut f: F) -> impl Iterator + '_ + pub fn drain_filter(&mut self, f: F) -> DrainFilter<'_, K, V, F> where - F: 'a + FnMut(&K, &mut V) -> bool, + F: FnMut(&K, &mut V) -> bool, { - // Here we only use `iter` as a temporary, preventing use-after-free - unsafe { - self.table.iter().filter_map(move |item| { - let &mut (ref key, ref mut value) = item.as_mut(); - if f(key, value) { - None - } else { - self.table.erase_no_drop(&item); - Some(item.read()) - } - }) + DrainFilter { + f, + iter: unsafe { self.table.iter() }, + table: &mut self.table, } } } @@ -1269,6 +1265,58 @@ impl Drain<'_, K, V> { } } +/// A draining iterator over entries of a `HashMap` which don't satisfy the predicate `f`. +/// +/// This `struct` is created by the [`drain_filter`] method on [`HashMap`]. See its +/// documentation for more. +/// +/// [`drain_filter`]: struct.HashMap.html#method.drain_filter +/// [`HashMap`]: struct.HashMap.html +pub struct DrainFilter<'a, K, V, F> +where + F: FnMut(&K, &mut V) -> bool, +{ + f: F, + iter: RawIter<(K, V)>, + table: &'a mut RawTable<(K, V)>, +} + +impl Drop for DrainFilter<'_, K, V, F> +where + F: FnMut(&K, &mut V) -> bool, +{ + fn drop(&mut self) { + unsafe { + while let Some(item) = self.iter.next() { + let &mut (ref key, ref mut value) = item.as_mut(); + if !(self.f)(key, value) { + self.table.erase_no_drop(&item); + item.drop(); + } + } + } + } +} + +impl Iterator for DrainFilter<'_, K, V, F> +where + F: FnMut(&K, &mut V) -> bool, +{ + type Item = (K, V); + fn next(&mut self) -> Option { + unsafe { + while let Some(item) = self.iter.next() { + let &mut (ref key, ref mut value) = item.as_mut(); + if !(self.f)(key, value) { + self.table.erase_no_drop(&item); + return Some(item.read()); + } + } + } + None + } +} + /// A mutable iterator over the values of a `HashMap`. /// /// This `struct` is created by the [`values_mut`] method on [`HashMap`]. See its @@ -3523,12 +3571,19 @@ mod test_map { #[test] fn test_drain_filter() { - let mut map: HashMap = (0..8).map(|x| (x, x * 10)).collect(); - let drained = map.drain_filter(|&k, _| k % 2 == 0); - let mut out = drained.collect::>(); - out.sort_unstable(); - assert_eq!(vec![(1, 10), (3, 30), (5, 50), (7, 70)], out); - assert_eq!(map.len(), 4); + { + let mut map: HashMap = (0..8).map(|x| (x, x * 10)).collect(); + let drained = map.drain_filter(|&k, _| k % 2 == 0); + let mut out = drained.collect::>(); + out.sort_unstable(); + assert_eq!(vec![(1, 10), (3, 30), (5, 50), (7, 70)], out); + assert_eq!(map.len(), 4); + } + { + let mut map: HashMap = (0..8).map(|x| (x, x * 10)).collect(); + drop(map.drain_filter(|&k, _| k % 2 == 0)); + assert_eq!(map.len(), 4); + } } #[test] From 0a65d86893a31b970ce02bcbe200ff917684bf90 Mon Sep 17 00:00:00 2001 From: julianknodt Date: Mon, 16 Dec 2019 10:52:49 -0800 Subject: [PATCH 4/4] Added drop guard --- src/map.rs | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/src/map.rs b/src/map.rs index f3e8c86e1b..5bb4a9d358 100644 --- a/src/map.rs +++ b/src/map.rs @@ -1281,20 +1281,28 @@ where table: &'a mut RawTable<(K, V)>, } -impl Drop for DrainFilter<'_, K, V, F> +impl<'a, K, V, F> Drop for DrainFilter<'a, K, V, F> where F: FnMut(&K, &mut V) -> bool, { fn drop(&mut self) { - unsafe { - while let Some(item) = self.iter.next() { - let &mut (ref key, ref mut value) = item.as_mut(); - if !(self.f)(key, value) { - self.table.erase_no_drop(&item); - item.drop(); - } + 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); + drop(item); + mem::forget(guard); + } } }