perf: Use for_each in Vec::extend

src/liballoc/benches/vec.rs

@@ -1,4 +1,7 @@
-use std::iter::{repeat, FromIterator};
+use std::{
+    hint::black_box,
+    iter::{repeat, FromIterator},
+};
 use test::Bencher;
 
 #[bench]
@@ -165,14 +168,53 @@ fn bench_from_iter_1000(b: &mut Bencher) {
 }
 
 fn do_bench_extend(b: &mut Bencher, dst_len: usize, src_len: usize) {
-    let dst: Vec<_> = FromIterator::from_iter(0..dst_len);
     let src: Vec<_> = FromIterator::from_iter(dst_len..dst_len + src_len);
+    do_bench_extend_iter(b, dst_len, src_len, src)
+}
+
+fn do_bench_extend_chain(b: &mut Bencher, dst_len: usize, src_len: usize) {
+    let src = Vec::from_iter(dst_len..dst_len + src_len);
+    let x1;
+    let x2;
+    let x3;
+    let x4;
+    let x5;
+
+    if src.len() / 5 == 0 {
+        x1 = &[][..];
+        x2 = &[][..];
+        x3 = &[][..];
+        x4 = &[][..];
+        x5 = &[][..];
+    } else {
+        let mut iter = src.chunks_exact(src.len() / 5);
+        x1 = iter.next().unwrap_or(&[][..]);
+        x2 = iter.next().unwrap_or(&[][..]);
+        x3 = iter.next().unwrap_or(&[][..]);
+        x4 = iter.next().unwrap_or(&[][..]);
+        x5 = iter.next().unwrap_or(&[][..]);
+    }
+    let src =
+        x1.iter().cloned().chain(x2.iter().cloned().chain(x3.iter().cloned())).chain(
+            Some(()).into_iter().flat_map(|()| x4.iter().cloned().chain(x5.iter().cloned())),
+        );
+    do_bench_extend_iter(b, dst_len, src_len, src)
+}
+
+#[inline(never)]
+fn do_bench_extend_iter(
+    b: &mut Bencher,
+    dst_len: usize,
+    src_len: usize,
+    src: impl IntoIterator<Item = usize> + Clone,
+) {
+    let dst: Vec<_> = FromIterator::from_iter(0..dst_len);
 
     b.bytes = src_len as u64;
 
     b.iter(|| {
-        let mut dst = dst.clone();
-        dst.extend(src.clone());
+        let mut dst = black_box(dst.clone());
+        dst.extend(black_box(src.clone()));
         assert_eq!(dst.len(), dst_len + src_len);
         assert!(dst.iter().enumerate().all(|(i, x)| i == *x));
     });
@@ -213,6 +255,41 @@ fn bench_extend_1000_1000(b: &mut Bencher) {
     do_bench_extend(b, 1000, 1000)
 }
 
+#[bench]
+fn bench_extend_chain_0000_0000(b: &mut Bencher) {
+    do_bench_extend_chain(b, 0, 0)
+}
+
+#[bench]
+fn bench_extend_chain_0000_0010(b: &mut Bencher) {
+    do_bench_extend_chain(b, 0, 10)
+}
+
+#[bench]
+fn bench_extend_chain_0000_0100(b: &mut Bencher) {
+    do_bench_extend_chain(b, 0, 100)
+}
+
+#[bench]
+fn bench_extend_chain_0000_1000(b: &mut Bencher) {
+    do_bench_extend_chain(b, 0, 1000)
+}
+
+#[bench]
+fn bench_extend_chain_0010_0010(b: &mut Bencher) {
+    do_bench_extend_chain(b, 10, 10)
+}
+
+#[bench]
+fn bench_extend_chain_0100_0100(b: &mut Bencher) {
+    do_bench_extend_chain(b, 100, 100)
+}
+
+#[bench]
+fn bench_extend_chain_1000_1000(b: &mut Bencher) {
+    do_bench_extend_chain(b, 1000, 1000)
+}
+
 fn do_bench_push_all(b: &mut Bencher, dst_len: usize, src_len: usize) {
     let dst: Vec<_> = FromIterator::from_iter(0..dst_len);
     let src: Vec<_> = FromIterator::from_iter(dst_len..dst_len + src_len);

src/liballoc/vec.rs

@@ -1559,8 +1559,8 @@ impl<T: Default> Vec<T> {
     #[unstable(feature = "vec_resize_default", issue = "41758")]
     #[rustc_deprecated(
         reason = "This is moving towards being removed in favor \
-        of `.resize_with(Default::default)`.  If you disagree, please comment \
-        in the tracking issue.",
+                  of `.resize_with(Default::default)`.  If you disagree, please comment \
+                  in the tracking issue.",
         since = "1.33.0"
     )]
     pub fn resize_default(&mut self, new_len: usize) {
@@ -2121,7 +2121,37 @@ where
     }
 }
 
+fn extend_fold<T>(self_: &mut Vec<T>) -> impl FnMut((), T) -> Result<(), T> + '_ {
+    move |(), element| {
+        if self_.len() < self_.capacity() {
+            unsafe {
+                self_.push_unchecked(element);
+            }
+            Ok(())
+        } else {
+            Err(element)
+        }
+    }
+}
+
 impl<T> Vec<T> {
+    unsafe fn push_unchecked(&mut self, value: T) {
+        let end = self.as_mut_ptr().add(self.len);
+        ptr::write(end, value);
+        // NB can't overflow since we would have had to alloc the address space
+        self.len += 1;
+    }
+
+    #[cold]
+    #[inline(never)]
+    fn extend_desugared_fallback<I: Iterator<Item = T>>(&mut self, iterator: &I, element: T) {
+        let (lower, _) = iterator.size_hint();
+        self.reserve(lower.saturating_add(1));
+        unsafe {
+            self.push_unchecked(element);
+        }
+    }
+
     fn extend_desugared<I: Iterator<Item = T>>(&mut self, mut iterator: I) {
         // This is the case for a general iterator.
         //
@@ -2130,17 +2160,10 @@ impl<T> Vec<T> {
         //      for item in iterator {
         //          self.push(item);
         //      }
-        while let Some(element) = iterator.next() {
-            let len = self.len();
-            if len == self.capacity() {
-                let (lower, _) = iterator.size_hint();
-                self.reserve(lower.saturating_add(1));
-            }
-            unsafe {
-                ptr::write(self.get_unchecked_mut(len), element);
-                // NB can't overflow since we would have had to alloc the address space
-                self.set_len(len + 1);
-            }
+        let (lower, _) = iterator.size_hint();
+        self.reserve(lower);
+        while let Err(element) = iterator.try_fold((), extend_fold(self)) {
+            self.extend_desugared_fallback(&iterator, element);
         }
     }