From 1b90dab9505cb07d1de24c1df1d0c073ff9cb397 Mon Sep 17 00:00:00 2001
From: leo60228 <leo@60228.dev>
Date: Tue, 12 May 2020 21:12:12 -0400
Subject: [PATCH] Make Iterator::unzip fast

Closes #72085

This consists of the following optimizations:

* Adds a `with_capacity` function to `Extend`. This definitely needs more thought if it's going to be stabilized, so I'm not writing an RFC yet. This takes off most of the performance gap.
* Optimizes `Vec`'s `Extend` implementation for the case where `size_hint` is 1. This shaves off the remaining performance gap.
---
 src/liballoc/lib.rs                 |  1 +
 src/liballoc/vec.rs                 | 15 ++++++++++++++-
 src/libcore/iter/traits/collect.rs  |  6 ++++++
 src/libcore/iter/traits/iterator.rs |  5 +++--
 4 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/src/liballoc/lib.rs b/src/liballoc/lib.rs
index 5365c9d01684..cc50d56563be 100644
--- a/src/liballoc/lib.rs
+++ b/src/liballoc/lib.rs
@@ -125,6 +125,7 @@
 #![feature(alloc_layout_extra)]
 #![feature(try_trait)]
 #![feature(associated_type_bounds)]
+#![feature(extend_with_capacity)]
 
 // Allow testing this library
 
diff --git a/src/liballoc/vec.rs b/src/liballoc/vec.rs
index cbfbf4d1cd32..2244b493b500 100644
--- a/src/liballoc/vec.rs
+++ b/src/liballoc/vec.rs
@@ -2048,6 +2048,11 @@ impl<T> Extend<T> for Vec<T> {
     fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
         <Self as SpecExtend<T, I::IntoIter>>::spec_extend(self, iter.into_iter())
     }
+
+    #[inline]
+    fn with_capacity(capacity: usize) -> Self {
+        Self::with_capacity(capacity)
+    }
 }
 
 // Specialization trait used for Vec::from_iter and Vec::extend
@@ -2097,7 +2102,7 @@ where
         vector
     }
 
-    default fn spec_extend(&mut self, iterator: I) {
+    default fn spec_extend(&mut self, mut iterator: I) {
         // This is the case for a TrustedLen iterator.
         let (low, high) = iterator.size_hint();
         if let Some(high_value) = high {
@@ -2109,6 +2114,14 @@ where
             );
         }
         if let Some(additional) = high {
+            if additional == 1 {
+                // work around inefficiencies in generic vec.extend(Some(x))
+                if let Some(x) = iterator.next() {
+                    self.push(x);
+                }
+                return;
+            }
+
             self.reserve(additional);
             unsafe {
                 let mut ptr = self.as_mut_ptr().add(self.len());
diff --git a/src/libcore/iter/traits/collect.rs b/src/libcore/iter/traits/collect.rs
index f21ab8dbc373..c08594b57e8b 100644
--- a/src/libcore/iter/traits/collect.rs
+++ b/src/libcore/iter/traits/collect.rs
@@ -341,6 +341,12 @@ pub trait Extend<A> {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     fn extend<T: IntoIterator<Item = A>>(&mut self, iter: T);
+
+    /// Creates this collection with a reserved capacity.
+    #[unstable(feature = "extend_with_capacity", issue = "none")]
+    fn with_capacity(_capacity: usize) -> Self where Self: Default {
+        Self::default()
+    }
 }
 
 #[stable(feature = "extend_for_unit", since = "1.28.0")]
diff --git a/src/libcore/iter/traits/iterator.rs b/src/libcore/iter/traits/iterator.rs
index 34ca79154b68..2bec04007baa 100644
--- a/src/libcore/iter/traits/iterator.rs
+++ b/src/libcore/iter/traits/iterator.rs
@@ -2672,8 +2672,9 @@ pub trait Iterator {
             }
         }
 
-        let mut ts: FromA = Default::default();
-        let mut us: FromB = Default::default();
+        let cap = self.size_hint().0.saturating_add(1);
+        let mut ts: FromA = Extend::with_capacity(cap);
+        let mut us: FromB = Extend::with_capacity(cap);
 
         self.for_each(extend(&mut ts, &mut us));