BTreeMap/BTreeSet drain & retain and share the now 7 similar algorithms

Author: ssomers

src/liballoc/benches/btree/map.rs

@@ -146,6 +146,36 @@ pub fn iter_100000(b: &mut Bencher) {
     bench_iter(b, 100000);
 }
 
+fn bench_iter_mut(b: &mut Bencher, size: i32) {
+    let mut map = BTreeMap::<i32, i32>::new();
+    let mut rng = thread_rng();
+
+    for _ in 0..size {
+        map.insert(rng.gen(), rng.gen());
+    }
+
+    b.iter(|| {
+        for kv in map.iter_mut() {
+            black_box(kv);
+        }
+    });
+}
+
+#[bench]
+pub fn iter_mut_20(b: &mut Bencher) {
+    bench_iter_mut(b, 20);
+}
+
+#[bench]
+pub fn iter_mut_1000(b: &mut Bencher) {
+    bench_iter_mut(b, 1000);
+}
+
+#[bench]
+pub fn iter_mut_100000(b: &mut Bencher) {
+    bench_iter_mut(b, 100000);
+}
+
 fn bench_first_and_last(b: &mut Bencher, size: i32) {
     let map: BTreeMap<_, _> = (0..size).map(|i| (i, i)).collect();
     b.iter(|| {

src/liballoc/benches/btree/set.rs

@@ -14,19 +14,13 @@ fn random(n: usize) -> BTreeSet<usize> {
 }
 
 fn neg(n: usize) -> BTreeSet<i32> {
-    let mut set = BTreeSet::new();
-    for i in -(n as i32)..=-1 {
-        set.insert(i);
-    }
+    let set: BTreeSet<i32> = (-(n as i32)..=-1).collect();
     assert_eq!(set.len(), n);
     set
 }
 
 fn pos(n: usize) -> BTreeSet<i32> {
-    let mut set = BTreeSet::new();
-    for i in 1..=(n as i32) {
-        set.insert(i);
-    }
+    let set: BTreeSet<i32> = (1..=(n as i32)).collect();
     assert_eq!(set.len(), n);
     set
 }
@@ -56,6 +50,54 @@ macro_rules! set_bench {
     };
 }
 
+const BUILD_SET_SIZE: usize = 100;
+
+#[bench]
+pub fn build_and_clear(b: &mut Bencher) {
+    b.iter(|| pos(BUILD_SET_SIZE).clear())
+}
+
+#[bench]
+pub fn build_and_drain(b: &mut Bencher) {
+    b.iter(|| pos(BUILD_SET_SIZE).drain().count())
+}
+
+#[bench]
+pub fn build_and_drop(b: &mut Bencher) {
+    b.iter(|| pos(BUILD_SET_SIZE))
+}
+
+#[bench]
+pub fn build_and_into_iter(b: &mut Bencher) {
+    b.iter(|| pos(BUILD_SET_SIZE).into_iter().count())
+}
+
+#[bench]
+pub fn build_and_pop_all(b: &mut Bencher) {
+    b.iter(|| {
+        let mut s = pos(BUILD_SET_SIZE);
+        while s.pop_first().is_some() {
+        }
+        s
+    });
+}
+
+#[bench]
+pub fn build_and_remove_all(b: &mut Bencher) {
+    b.iter(|| {
+        let mut s = pos(BUILD_SET_SIZE);
+        while let Some(elt) = s.iter().copied().next() {
+            s.remove(&elt);
+        }
+        s
+    });
+}
+
+#[bench]
+pub fn build_and_retain_nothing(b: &mut Bencher) {
+    b.iter(|| pos(BUILD_SET_SIZE).retain(|_| false))
+}
+
 set_bench! {intersection_100_neg_vs_100_pos, intersection, count, [neg(100), pos(100)]}
 set_bench! {intersection_100_neg_vs_10k_pos, intersection, count, [neg(100), pos(10_000)]}
 set_bench! {intersection_100_pos_vs_100_neg, intersection, count, [pos(100), neg(100)]}

src/liballoc/benches/lib.rs

@@ -1,3 +1,4 @@
+#![feature(btree_drain_retain)]
 #![feature(map_first_last)]
 #![feature(repr_simd)]
 #![feature(test)]