Auto merge of #43245 - Gankro:drain-filter, r=sfackler

Add Vec::drain_filter

This implements the API proposed in #43244.

So I spent like half a day figuring out how to implement this in some awesome super-optimized unsafe way, which had me very confident this was worth putting into the stdlib.

Then I looked at the impl for `retain`, and was like "oh dang". I compared the two and they basically ended up being the same speed. And the `retain` impl probably translates to DoubleEndedIter a lot more cleanly if we ever want that.

So now I'm not totally confident this needs to go in the stdlib, but I've got two implementations and an amazingly robust test suite, so I figured I might as well toss it over the fence for discussion.

Author: bors

src/liballoc/tests/lib.rs

@@ -16,6 +16,7 @@
 #![feature(inclusive_range_syntax)]
 #![feature(collection_placement)]
 #![feature(const_fn)]
+#![feature(drain_filter)]
 #![feature(exact_size_is_empty)]
 #![feature(iterator_step_by)]
 #![feature(pattern)]

src/liballoc/tests/vec.rs

@@ -801,3 +801,170 @@ fn overaligned_allocations() {
         assert!(v.as_ptr() as usize & 0xff == 0);
     }
 }
+
+#[test]
+fn drain_filter_empty() {
+    let mut vec: Vec<i32> = vec![];
+
+    {
+        let mut iter = vec.drain_filter(|_| true);
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+        assert_eq!(iter.next(), None);
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+        assert_eq!(iter.next(), None);
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+    }
+    assert_eq!(vec.len(), 0);
+    assert_eq!(vec, vec![]);
+}
+
+#[test]
+fn drain_filter_zst() {
+    let mut vec = vec![(), (), (), (), ()];
+    let initial_len = vec.len();
+    let mut count = 0;
+    {
+        let mut iter = vec.drain_filter(|_| true);
+        assert_eq!(iter.size_hint(), (0, Some(initial_len)));
+        while let Some(_) = iter.next() {
+            count += 1;
+            assert_eq!(iter.size_hint(), (0, Some(initial_len - count)));
+        }
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+        assert_eq!(iter.next(), None);
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+    }
+
+    assert_eq!(count, initial_len);
+    assert_eq!(vec.len(), 0);
+    assert_eq!(vec, vec![]);
+}
+
+#[test]
+fn drain_filter_false() {
+    let mut vec = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+
+    let initial_len = vec.len();
+    let mut count = 0;
+    {
+        let mut iter = vec.drain_filter(|_| false);
+        assert_eq!(iter.size_hint(), (0, Some(initial_len)));
+        for _ in iter.by_ref() {
+            count += 1;
+        }
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+        assert_eq!(iter.next(), None);
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+    }
+
+    assert_eq!(count, 0);
+    assert_eq!(vec.len(), initial_len);
+    assert_eq!(vec, vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+}
+
+#[test]
+fn drain_filter_true() {
+    let mut vec = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+
+    let initial_len = vec.len();
+    let mut count = 0;
+    {
+        let mut iter = vec.drain_filter(|_| true);
+        assert_eq!(iter.size_hint(), (0, Some(initial_len)));
+        while let Some(_) = iter.next() {
+            count += 1;
+            assert_eq!(iter.size_hint(), (0, Some(initial_len - count)));
+        }
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+        assert_eq!(iter.next(), None);
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+    }
+
+    assert_eq!(count, initial_len);
+    assert_eq!(vec.len(), 0);
+    assert_eq!(vec, vec![]);
+}
+
+#[test]
+fn drain_filter_complex() {
+
+    {   //                [+xxx++++++xxxxx++++x+x++]
+        let mut vec = vec![1,
+                           2, 4, 6,
+                           7, 9, 11, 13, 15, 17,
+                           18, 20, 22, 24, 26,
+                           27, 29, 31, 33,
+                           34,
+                           35,
+                           36,
+                           37, 39];
+
+        let removed = vec.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
+        assert_eq!(removed.len(), 10);
+        assert_eq!(removed, vec![2, 4, 6, 18, 20, 22, 24, 26, 34, 36]);
+
+        assert_eq!(vec.len(), 14);
+        assert_eq!(vec, vec![1, 7, 9, 11, 13, 15, 17, 27, 29, 31, 33, 35, 37, 39]);
+    }
+
+    {   //                [xxx++++++xxxxx++++x+x++]
+        let mut vec = vec![2, 4, 6,
+                           7, 9, 11, 13, 15, 17,
+                           18, 20, 22, 24, 26,
+                           27, 29, 31, 33,
+                           34,
+                           35,
+                           36,
+                           37, 39];
+
+        let removed = vec.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
+        assert_eq!(removed.len(), 10);
+        assert_eq!(removed, vec![2, 4, 6, 18, 20, 22, 24, 26, 34, 36]);
+
+        assert_eq!(vec.len(), 13);
+        assert_eq!(vec, vec![7, 9, 11, 13, 15, 17, 27, 29, 31, 33, 35, 37, 39]);
+    }
+
+    {   //                [xxx++++++xxxxx++++x+x]
+        let mut vec = vec![2, 4, 6,
+                           7, 9, 11, 13, 15, 17,
+                           18, 20, 22, 24, 26,
+                           27, 29, 31, 33,
+                           34,
+                           35,
+                           36];
+
+        let removed = vec.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
+        assert_eq!(removed.len(), 10);
+        assert_eq!(removed, vec![2, 4, 6, 18, 20, 22, 24, 26, 34, 36]);
+
+        assert_eq!(vec.len(), 11);
+        assert_eq!(vec, vec![7, 9, 11, 13, 15, 17, 27, 29, 31, 33, 35]);
+    }
+
+    {   //                [xxxxxxxxxx+++++++++++]
+        let mut vec = vec![2, 4, 6, 8, 10, 12, 14, 16, 18, 20,
+                           1, 3, 5, 7, 9, 11, 13, 15, 17, 19];
+
+        let removed = vec.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
+        assert_eq!(removed.len(), 10);
+        assert_eq!(removed, vec![2, 4, 6, 8, 10, 12, 14, 16, 18, 20]);
+
+        assert_eq!(vec.len(), 10);
+        assert_eq!(vec, vec![1, 3, 5, 7, 9, 11, 13, 15, 17, 19]);
+    }
+
+    {   //                [+++++++++++xxxxxxxxxx]
+        let mut vec = vec![1, 3, 5, 7, 9, 11, 13, 15, 17, 19,
+                           2, 4, 6, 8, 10, 12, 14, 16, 18, 20];
+
+        let removed = vec.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
+        assert_eq!(removed.len(), 10);
+        assert_eq!(removed, vec![2, 4, 6, 8, 10, 12, 14, 16, 18, 20]);
+
+        assert_eq!(vec.len(), 10);
+        assert_eq!(vec, vec![1, 3, 5, 7, 9, 11, 13, 15, 17, 19]);
+    }
+}
+
+

src/liballoc/vec.rs

@@ -1960,6 +1960,65 @@ impl<T> Vec<T> {
         }
     }
 
+    /// Creates an iterator which uses a closure to determine if an element should be removed.
+    ///
+    /// If the closure returns true, then the element is removed and yielded.
+    /// If the closure returns false, it will try again, and call the closure
+    /// on the next element, seeing if it passes the test.
+    ///
+    /// Using this method is equivalent to the following code:
+    ///
+    /// ```
+    /// # let some_predicate = |x: &mut i32| { *x == 2 };
+    /// # let mut vec = vec![1, 2, 3, 4, 5];
+    /// let mut i = 0;
+    /// while i != vec.len() {
+    ///     if some_predicate(&mut vec[i]) {
+    ///         let val = vec.remove(i);
+    ///         // your code here
+    ///     }
+    ///     i += 1;
+    /// }
+    /// ```
+    ///
+    /// But `drain_filter` is easier to use. `drain_filter` is also more efficient,
+    /// because it can backshift the elements of the array in bulk.
+    ///
+    /// Note that `drain_filter` also lets you mutate every element in the filter closure,
+    /// regardless of whether you choose to keep or remove it.
+    ///
+    ///
+    /// # Examples
+    ///
+    /// Splitting an array into evens and odds, reusing the original allocation:
+    ///
+    /// ```
+    /// #![feature(drain_filter)]
+    /// let mut numbers = vec![1, 2, 3, 4, 5, 6, 8, 9, 11, 13, 14, 15];
+    ///
+    /// let evens = numbers.drain_filter(|x| *x % 2 == 0).collect::<Vec<_>>();
+    /// let odds = numbers;
+    ///
+    /// assert_eq!(evens, vec![2, 4, 6, 8, 14]);
+    /// assert_eq!(odds, vec![1, 3, 5, 9, 11, 13, 15]);
+    /// ```
+    #[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
+    pub fn drain_filter<F>(&mut self, filter: F) -> DrainFilter<T, F>
+        where F: FnMut(&mut T) -> bool,
+    {
+        let old_len = self.len();
+
+        // Guard against us getting leaked (leak amplification)
+        unsafe { self.set_len(0); }
+
+        DrainFilter {
+            vec: self,
+            idx: 0,
+            del: 0,
+            old_len,
+            pred: filter,
+        }
+    }
 }
 
 /// Extend implementation that copies elements out of references before pushing them onto the Vec.
@@ -2602,3 +2661,57 @@ impl<'a, T> Drain<'a, T> {
         self.tail_start = new_tail_start;
     }
 }
+
+/// An iterator produced by calling `drain_filter` on Vec.
+#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
+#[derive(Debug)]
+pub struct DrainFilter<'a, T: 'a, F>
+    where F: FnMut(&mut T) -> bool,
+{
+    vec: &'a mut Vec<T>,
+    idx: usize,
+    del: usize,
+    old_len: usize,
+    pred: F,
+}
+
+#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
+impl<'a, T, F> Iterator for DrainFilter<'a, T, F>
+    where F: FnMut(&mut T) -> bool,
+{
+    type Item = T;
+
+    fn next(&mut self) -> Option<T> {
+        unsafe {
+            while self.idx != self.old_len {
+                let i = self.idx;
+                self.idx += 1;
+                let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len);
+                if (self.pred)(&mut v[i]) {
+                    self.del += 1;
+                    return Some(ptr::read(&v[i]));
+                } else if self.del > 0 {
+                    v.swap(i - self.del, i);
+                }
+            }
+            None
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (0, Some(self.old_len - self.idx))
+    }
+}
+
+#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
+impl<'a, T, F> Drop for DrainFilter<'a, T, F>
+    where F: FnMut(&mut T) -> bool,
+{
+    fn drop(&mut self) {
+        for _ in self.by_ref() { }
+
+        unsafe {
+            self.vec.set_len(self.old_len - self.del);
+        }
+    }
+}