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.

2 files changed, 168 insertions, 0 deletions

diff --git a/src/liballoc/tests/lib.rs b/src/liballoc/tests/lib.rs
index 34b0a137b10..86309dd87de 100644
--- a/src/liballoc/tests/lib.rs
+++ b/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)]
diff --git a/src/liballoc/tests/vec.rs b/src/liballoc/tests/vec.rs
index 17f1229c206..670ea8089fc 100644
--- a/src/liballoc/tests/vec.rs
+++ b/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]);
+    }
+}
+
+