Rollup merge of #72162 - cuviper:extend_one, r=Mark-Simulacrum

Add Extend::{extend_one,extend_reserve}

This adds new optional methods on `Extend`: `extend_one` add a single
element to the collection, and `extend_reserve` pre-allocates space for
the predicted number of incoming elements. These are used in `Iterator`
for `partition` and `unzip` as they shuffle elements one-at-a-time into
their respective collections.

2 files changed, 26 insertions, 5 deletions

diff --git a/src/libcore/iter/traits/collect.rs b/src/libcore/iter/traits/collect.rs
index f21ab8dbc37..9d20022b6ed 100644
--- a/src/libcore/iter/traits/collect.rs
+++ b/src/libcore/iter/traits/collect.rs
@@ -322,7 +322,7 @@ impl<I: Iterator> IntoIterator for I {
 pub trait Extend<A> {
     /// Extends a collection with the contents of an iterator.
     ///
-    /// As this is the only method for this trait, the [trait-level] docs
+    /// As this is the only required method for this trait, the [trait-level] docs
     /// contain more details.
     ///
     /// [trait-level]: trait.Extend.html
@@ -341,6 +341,20 @@ pub trait Extend<A> {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     fn extend<T: IntoIterator<Item = A>>(&mut self, iter: T);
+
+    /// Extends a collection with exactly one element.
+    #[unstable(feature = "extend_one", issue = "72631")]
+    fn extend_one(&mut self, item: A) {
+        self.extend(Some(item));
+    }
+
+    /// Reserves capacity in a collection for the given number of additional elements.
+    ///
+    /// The default implementation does nothing.
+    #[unstable(feature = "extend_one", issue = "72631")]
+    fn extend_reserve(&mut self, additional: usize) {
+        let _ = additional;
+    }
 }
 
 #[stable(feature = "extend_for_unit", since = "1.28.0")]
@@ -348,4 +362,5 @@ impl Extend<()> for () {
     fn extend<T: IntoIterator<Item = ()>>(&mut self, iter: T) {
         iter.into_iter().for_each(drop)
     }
+    fn extend_one(&mut self, _item: ()) {}
 }
diff --git a/src/libcore/iter/traits/iterator.rs b/src/libcore/iter/traits/iterator.rs
index 7f081f732fd..a10b34d931d 100644
--- a/src/libcore/iter/traits/iterator.rs
+++ b/src/libcore/iter/traits/iterator.rs
@@ -1711,9 +1711,9 @@ pub trait Iterator {
         ) -> impl FnMut((), T) + 'a {
             move |(), x| {
                 if f(&x) {
-                    left.extend(Some(x));
+                    left.extend_one(x);
                 } else {
-                    right.extend(Some(x));
+                    right.extend_one(x);
                 }
             }
         }
@@ -2686,14 +2686,20 @@ pub trait Iterator {
             us: &'a mut impl Extend<B>,
         ) -> impl FnMut((), (A, B)) + 'a {
             move |(), (t, u)| {
-                ts.extend(Some(t));
-                us.extend(Some(u));
+                ts.extend_one(t);
+                us.extend_one(u);
             }
         }
 
         let mut ts: FromA = Default::default();
         let mut us: FromB = Default::default();
 
+        let (lower_bound, _) = self.size_hint();
+        if lower_bound > 0 {
+            ts.extend_reserve(lower_bound);
+            us.extend_reserve(lower_bound);
+        }
+
         self.fold((), extend(&mut ts, &mut us));
 
         (ts, us)