Auto merge of #66383 - the8472:in-place-iter-collect, r=<try>

Extend vec::FromIterator specialization to collect in-place for some iterator pipelines

This extends the existing `vec.into_iter().collect::<Vec>()` specialization to cover more cases, including iterator pipelines that flow from `IntoIter` through `Map`, `Filter`, `Zip` and `Enumerate` and several others.

[Benchmark results](https://gist.github.com/the8472/6d999b2d08a2bedf3b93f12112f96e2f) with defaults and opt-level = 3, codegen-units=1.

The specialization should be conservative, it is only applied when we know the iterator will fit into the source allocation.

Additionally some existing specializations now applied in a few more code-paths.

Author: bors

src/liballoc/benches/vec.rs

@@ -1,12 +1,14 @@
+use rand::prelude::*;
 use std::iter::{repeat, FromIterator};
-use test::Bencher;
+use test::{black_box, Bencher};
 
 #[bench]
 fn bench_new(b: &mut Bencher) {
     b.iter(|| {
         let v: Vec<u32> = Vec::new();
         assert_eq!(v.len(), 0);
         assert_eq!(v.capacity(), 0);
+        v
     })
 }
 
@@ -17,6 +19,7 @@ fn do_bench_with_capacity(b: &mut Bencher, src_len: usize) {
         let v: Vec<u32> = Vec::with_capacity(src_len);
         assert_eq!(v.len(), 0);
         assert_eq!(v.capacity(), src_len);
+        v
     })
 }
 
@@ -47,6 +50,7 @@ fn do_bench_from_fn(b: &mut Bencher, src_len: usize) {
         let dst = (0..src_len).collect::<Vec<_>>();
         assert_eq!(dst.len(), src_len);
         assert!(dst.iter().enumerate().all(|(i, x)| i == *x));
+        dst
     })
 }
 
@@ -77,6 +81,7 @@ fn do_bench_from_elem(b: &mut Bencher, src_len: usize) {
         let dst: Vec<usize> = repeat(5).take(src_len).collect();
         assert_eq!(dst.len(), src_len);
         assert!(dst.iter().all(|x| *x == 5));
+        dst
     })
 }
 
@@ -109,6 +114,7 @@ fn do_bench_from_slice(b: &mut Bencher, src_len: usize) {
         let dst = src.clone()[..].to_vec();
         assert_eq!(dst.len(), src_len);
         assert!(dst.iter().enumerate().all(|(i, x)| i == *x));
+        dst
     });
 }
 
@@ -141,6 +147,7 @@ fn do_bench_from_iter(b: &mut Bencher, src_len: usize) {
         let dst: Vec<_> = FromIterator::from_iter(src.clone());
         assert_eq!(dst.len(), src_len);
         assert!(dst.iter().enumerate().all(|(i, x)| i == *x));
+        dst
     });
 }
 
@@ -175,6 +182,7 @@ fn do_bench_extend(b: &mut Bencher, dst_len: usize, src_len: usize) {
         dst.extend(src.clone());
         assert_eq!(dst.len(), dst_len + src_len);
         assert!(dst.iter().enumerate().all(|(i, x)| i == *x));
+        dst
     });
 }
 
@@ -224,9 +232,24 @@ fn do_bench_push_all(b: &mut Bencher, dst_len: usize, src_len: usize) {
         dst.extend_from_slice(&src);
         assert_eq!(dst.len(), dst_len + src_len);
         assert!(dst.iter().enumerate().all(|(i, x)| i == *x));
+        dst
     });
 }
 
+#[bench]
+fn bench_extend_recycle(b: &mut Bencher) {
+    let mut data = vec![0; 1000];
+
+    b.iter(|| {
+        let tmp = std::mem::replace(&mut data, Vec::new());
+        let mut to_extend = black_box(Vec::new());
+        to_extend.extend(tmp.into_iter());
+        std::mem::replace(&mut data, black_box(to_extend));
+    });
+
+    black_box(data);
+}
+
 #[bench]
 fn bench_push_all_0000_0000(b: &mut Bencher) {
     do_bench_push_all(b, 0, 0)
@@ -273,6 +296,7 @@ fn do_bench_push_all_move(b: &mut Bencher, dst_len: usize, src_len: usize) {
         dst.extend(src.clone());
         assert_eq!(dst.len(), dst_len + src_len);
         assert!(dst.iter().enumerate().all(|(i, x)| i == *x));
+        dst
     });
 }
 
@@ -320,6 +344,7 @@ fn do_bench_clone(b: &mut Bencher, src_len: usize) {
         let dst = src.clone();
         assert_eq!(dst.len(), src_len);
         assert!(dst.iter().enumerate().all(|(i, x)| i == *x));
+        dst
     });
 }
 
@@ -354,10 +379,10 @@ fn do_bench_clone_from(b: &mut Bencher, times: usize, dst_len: usize, src_len: u
 
         for _ in 0..times {
             dst.clone_from(&src);
-
             assert_eq!(dst.len(), src_len);
             assert!(dst.iter().enumerate().all(|(i, x)| dst_len + i == *x));
         }
+        dst
     });
 }
 
@@ -480,3 +505,223 @@ fn bench_clone_from_10_0100_0010(b: &mut Bencher) {
 fn bench_clone_from_10_1000_0100(b: &mut Bencher) {
     do_bench_clone_from(b, 10, 1000, 100)
 }
+
+macro_rules! bench_in_place {
+    (
+        $($fname:ident, $type:ty , $count:expr, $init: expr);*
+    ) => {
+        $(
+            #[bench]
+            fn $fname(b: &mut Bencher) {
+                b.iter(|| {
+                    let src: Vec<$type> = black_box(vec![$init; $count]);
+                    let mut sink = src.into_iter()
+                        .enumerate()
+                        .map(|(idx, e)| { (idx as $type) ^ e }).collect::<Vec<$type>>();
+                    black_box(sink.as_mut_ptr())
+                });
+            }
+        )+
+    };
+}
+
+bench_in_place![
+    bench_in_place_xxu8_i0_0010,     u8,     10, 0;
+    bench_in_place_xxu8_i0_0100,     u8,    100, 0;
+    bench_in_place_xxu8_i0_1000,     u8,   1000, 0;
+    bench_in_place_xxu8_i1_0010,     u8,     10, 1;
+    bench_in_place_xxu8_i1_0100,     u8,    100, 1;
+    bench_in_place_xxu8_i1_1000,     u8,   1000, 1;
+    bench_in_place_xu32_i0_0010,    u32,     10, 0;
+    bench_in_place_xu32_i0_0100,    u32,    100, 0;
+    bench_in_place_xu32_i0_1000,    u32,   1000, 0;
+    bench_in_place_xu32_i1_0010,    u32,     10, 1;
+    bench_in_place_xu32_i1_0100,    u32,    100, 1;
+    bench_in_place_xu32_i1_1000,    u32,   1000, 1;
+    bench_in_place_u128_i0_0010,   u128,     10, 0;
+    bench_in_place_u128_i0_0100,   u128,    100, 0;
+    bench_in_place_u128_i0_1000,   u128,   1000, 0;
+    bench_in_place_u128_i1_0010,   u128,     10, 1;
+    bench_in_place_u128_i1_0100,   u128,    100, 1;
+    bench_in_place_u128_i1_1000,   u128,   1000, 1
+];
+
+#[bench]
+fn bench_in_place_recycle(b: &mut test::Bencher) {
+    let mut data = vec![0; 1000];
+
+    b.iter(|| {
+        let tmp = std::mem::replace(&mut data, Vec::new());
+        std::mem::replace(
+            &mut data,
+            black_box(
+                tmp.into_iter()
+                    .enumerate()
+                    .map(|(idx, e)| idx.wrapping_add(e))
+                    .fuse()
+                    .peekable()
+                    .collect::<Vec<usize>>(),
+            ),
+        );
+    });
+}
+
+#[bench]
+fn bench_in_place_zip_recycle(b: &mut test::Bencher) {
+    let mut data = vec![0u8; 1000];
+    let mut rng = rand::thread_rng();
+    let mut subst = vec![0u8; 1000];
+    rng.fill_bytes(&mut subst[..]);
+
+    b.iter(|| {
+        let tmp = std::mem::replace(&mut data, Vec::new());
+        let mangled = tmp
+            .into_iter()
+            .zip(subst.iter().copied())
+            .enumerate()
+            .map(|(i, (d, s))| d.wrapping_add(i as u8) ^ s)
+            .collect::<Vec<_>>();
+        assert_eq!(mangled.len(), 1000);
+        std::mem::replace(&mut data, black_box(mangled));
+    });
+}
+
+#[bench]
+fn bench_in_place_zip_iter_mut(b: &mut test::Bencher) {
+    let mut data = vec![0u8; 256];
+    let mut rng = rand::thread_rng();
+    let mut subst = vec![0u8; 1000];
+    rng.fill_bytes(&mut subst[..]);
+
+    b.iter(|| {
+        data.iter_mut().enumerate().for_each(|(i, d)| {
+            *d = d.wrapping_add(i as u8) ^ subst[i];
+        });
+    });
+
+    black_box(data);
+}
+
+#[derive(Clone)]
+struct Droppable(usize);
+
+impl Drop for Droppable {
+    fn drop(&mut self) {
+        black_box(self);
+    }
+}
+
+#[bench]
+fn bench_in_place_collect_droppable(b: &mut test::Bencher) {
+    let v: Vec<Droppable> = std::iter::repeat_with(|| Droppable(0)).take(1000).collect();
+    b.iter(|| {
+        v.clone()
+            .into_iter()
+            .skip(100)
+            .enumerate()
+            .map(|(i, e)| Droppable(i ^ e.0))
+            .collect::<Vec<_>>()
+    })
+}
+
+#[bench]
+fn bench_chain_collect(b: &mut test::Bencher) {
+    let data = black_box([0; LEN]);
+    b.iter(|| data.iter().cloned().chain([1].iter().cloned()).collect::<Vec<_>>());
+}
+
+#[bench]
+fn bench_chain_chain_collect(b: &mut test::Bencher) {
+    let data = black_box([0; LEN]);
+    b.iter(|| {
+        data.iter()
+            .cloned()
+            .chain([1].iter().cloned())
+            .chain([2].iter().cloned())
+            .collect::<Vec<_>>()
+    });
+}
+
+#[bench]
+fn bench_nest_chain_chain_collect(b: &mut test::Bencher) {
+    let data = black_box([0; LEN]);
+    b.iter(|| {
+        data.iter().cloned().chain([1].iter().chain([2].iter()).cloned()).collect::<Vec<_>>()
+    });
+}
+
+pub fn example_plain_slow(l: &[u32]) -> Vec<u32> {
+    let mut result = Vec::with_capacity(l.len());
+    result.extend(l.iter().rev());
+    result
+}
+
+pub fn map_fast(l: &[(u32, u32)]) -> Vec<u32> {
+    let mut result = Vec::with_capacity(l.len());
+    for i in 0..l.len() {
+        unsafe {
+            *result.get_unchecked_mut(i) = l[i].0;
+            result.set_len(i);
+        }
+    }
+    result
+}
+
+const LEN: usize = 16384;
+
+#[bench]
+fn bench_range_map_collect(b: &mut test::Bencher) {
+    b.iter(|| (0..LEN).map(|_| u32::default()).collect::<Vec<_>>());
+}
+
+#[bench]
+fn bench_chain_extend_ref(b: &mut test::Bencher) {
+    let data = black_box([0; LEN]);
+    b.iter(|| {
+        let mut v = Vec::<u32>::with_capacity(data.len() + 1);
+        v.extend(data.iter().chain([1].iter()));
+        v
+    });
+}
+
+#[bench]
+fn bench_chain_extend_value(b: &mut test::Bencher) {
+    let data = black_box([0; LEN]);
+    b.iter(|| {
+        let mut v = Vec::<u32>::with_capacity(data.len() + 1);
+        v.extend(data.iter().cloned().chain(Some(1)));
+        v
+    });
+}
+
+#[bench]
+fn bench_rev_1(b: &mut test::Bencher) {
+    let data = black_box([0; LEN]);
+    b.iter(|| {
+        let mut v = Vec::<u32>::new();
+        v.extend(data.iter().rev());
+        v
+    });
+}
+
+#[bench]
+fn bench_rev_2(b: &mut test::Bencher) {
+    let data = black_box([0; LEN]);
+    b.iter(|| example_plain_slow(&data));
+}
+
+#[bench]
+fn bench_map_regular(b: &mut test::Bencher) {
+    let data = black_box([(0, 0); LEN]);
+    b.iter(|| {
+        let mut v = Vec::<u32>::new();
+        v.extend(data.iter().map(|t| t.1));
+        v
+    });
+}
+
+#[bench]
+fn bench_map_fast(b: &mut test::Bencher) {
+    let data = black_box([(0, 0); LEN]);
+    b.iter(|| map_fast(&data));
+}

src/liballoc/collections/binary_heap.rs

@@ -146,13 +146,13 @@
 #![stable(feature = "rust1", since = "1.0.0")]
 
 use core::fmt;
-use core::iter::{FromIterator, FusedIterator, TrustedLen};
+use core::iter::{FromIterator, FusedIterator, InPlaceIterable, SourceIter, TrustedLen};
 use core::mem::{size_of, swap, ManuallyDrop};
 use core::ops::{Deref, DerefMut};
 use core::ptr;
 
 use crate::slice;
-use crate::vec::{self, Vec};
+use crate::vec::{self, AsIntoIter, Vec};
 
 use super::SpecExtend;
 
@@ -1145,6 +1145,27 @@ impl<T> ExactSizeIterator for IntoIter<T> {
 #[stable(feature = "fused", since = "1.26.0")]
 impl<T> FusedIterator for IntoIter<T> {}
 
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<T> SourceIter for IntoIter<T> {
+    type Source = IntoIter<T>;
+
+    #[inline]
+    unsafe fn as_inner(&mut self) -> &mut Self::Source {
+        self
+    }
+}
+
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<I> InPlaceIterable for IntoIter<I> {}
+
+impl<I> AsIntoIter for IntoIter<I> {
+    type Item = I;
+
+    fn as_into_iter(&mut self) -> &mut vec::IntoIter<Self::Item> {
+        &mut self.iter
+    }
+}
+
 #[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")]
 #[derive(Clone, Debug)]
 pub struct IntoIterSorted<T> {

src/liballoc/lib.rs

@@ -121,6 +121,9 @@
 #![feature(alloc_layout_extra)]
 #![feature(try_trait)]
 #![feature(associated_type_bounds)]
+#![feature(inplace_iteration)]
+#![feature(type_alias_impl_trait)]
+#![feature(never_type)]
 
 // Allow testing this library