BTreeMap/BTreeSet: implement and test drain_filter

2 files changed, 15 insertions, 6 deletions

diff --git a/src/liballoc/benches/btree/set.rs b/src/liballoc/benches/btree/set.rs
index c4712e2132f..2518506b9b5 100644
--- a/src/liballoc/benches/btree/set.rs
+++ b/src/liballoc/benches/btree/set.rs
@@ -63,13 +63,17 @@ pub fn clone_100_and_clear(b: &mut Bencher) {
 }
 
 #[bench]
+pub fn clone_100_and_drain_all(b: &mut Bencher) {
+    let src = pos(100);
+    b.iter(|| src.clone().drain_filter(|_| true).count())
+}
+
+#[bench]
 pub fn clone_100_and_drain_half(b: &mut Bencher) {
     let src = pos(100);
     b.iter(|| {
         let mut set = src.clone();
-        for i in set.iter().copied().filter(|i| i % 2 == 0).collect::<Vec<_>>() {
-            set.remove(&i);
-        }
+        assert_eq!(set.drain_filter(|i| i % 2 == 0).count(), 100 / 2);
         assert_eq!(set.len(), 100 / 2);
     })
 }
@@ -128,13 +132,17 @@ pub fn clone_10k_and_clear(b: &mut Bencher) {
 }
 
 #[bench]
+pub fn clone_10k_and_drain_all(b: &mut Bencher) {
+    let src = pos(10_000);
+    b.iter(|| src.clone().drain_filter(|_| true).count())
+}
+
+#[bench]
 pub fn clone_10k_and_drain_half(b: &mut Bencher) {
     let src = pos(10_000);
     b.iter(|| {
         let mut set = src.clone();
-        for i in set.iter().copied().filter(|i| i % 2 == 0).collect::<Vec<_>>() {
-            set.remove(&i);
-        }
+        assert_eq!(set.drain_filter(|i| i % 2 == 0).count(), 10_000 / 2);
         assert_eq!(set.len(), 10_000 / 2);
     })
 }
diff --git a/src/liballoc/benches/lib.rs b/src/liballoc/benches/lib.rs
index 951477a24c8..f31717d9fd5 100644
--- a/src/liballoc/benches/lib.rs
+++ b/src/liballoc/benches/lib.rs
@@ -1,3 +1,4 @@
+#![feature(btree_drain_filter)]
 #![feature(map_first_last)]
 #![feature(repr_simd)]
 #![feature(test)]