Rollup merge of rust-lang#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.

src/liballoc/collections/binary_heap.rs

@@ -1376,6 +1376,16 @@ impl<T: Ord> Extend<T> for BinaryHeap<T> {
     fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
         <Self as SpecExtend<I>>::spec_extend(self, iter);
     }
+
+    #[inline]
+    fn extend_one(&mut self, item: T) {
+        self.push(item);
+    }
+
+    #[inline]
+    fn extend_reserve(&mut self, additional: usize) {
+        self.reserve(additional);
+    }
 }
 
 impl<T: Ord, I: IntoIterator<Item = T>> SpecExtend<I> for BinaryHeap<T> {
@@ -1406,4 +1416,14 @@ impl<'a, T: 'a + Ord + Copy> Extend<&'a T> for BinaryHeap<T> {
     fn extend<I: IntoIterator<Item = &'a T>>(&mut self, iter: I) {
         self.extend(iter.into_iter().cloned());
     }
+
+    #[inline]
+    fn extend_one(&mut self, &item: &'a T) {
+        self.push(item);
+    }
+
+    #[inline]
+    fn extend_reserve(&mut self, additional: usize) {
+        self.reserve(additional);
+    }
 }

src/liballoc/collections/btree/map.rs

@@ -1901,13 +1901,23 @@ impl<K: Ord, V> Extend<(K, V)> for BTreeMap<K, V> {
             self.insert(k, v);
         });
     }
+
+    #[inline]
+    fn extend_one(&mut self, (k, v): (K, V)) {
+        self.insert(k, v);
+    }
 }
 
 #[stable(feature = "extend_ref", since = "1.2.0")]
 impl<'a, K: Ord + Copy, V: Copy> Extend<(&'a K, &'a V)> for BTreeMap<K, V> {
     fn extend<I: IntoIterator<Item = (&'a K, &'a V)>>(&mut self, iter: I) {
         self.extend(iter.into_iter().map(|(&key, &value)| (key, value)));
     }
+
+    #[inline]
+    fn extend_one(&mut self, (&k, &v): (&'a K, &'a V)) {
+        self.insert(k, v);
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]

src/liballoc/collections/btree/set.rs

@@ -1152,13 +1152,23 @@ impl<T: Ord> Extend<T> for BTreeSet<T> {
             self.insert(elem);
         });
     }
+
+    #[inline]
+    fn extend_one(&mut self, elem: T) {
+        self.insert(elem);
+    }
 }
 
 #[stable(feature = "extend_ref", since = "1.2.0")]
 impl<'a, T: 'a + Ord + Copy> Extend<&'a T> for BTreeSet<T> {
     fn extend<I: IntoIterator<Item = &'a T>>(&mut self, iter: I) {
         self.extend(iter.into_iter().cloned());
     }
+
+    #[inline]
+    fn extend_one(&mut self, &elem: &'a T) {
+        self.insert(elem);
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]

src/liballoc/collections/linked_list.rs

@@ -1748,6 +1748,11 @@ impl<T> Extend<T> for LinkedList<T> {
     fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
         <Self as SpecExtend<I>>::spec_extend(self, iter);
     }
+
+    #[inline]
+    fn extend_one(&mut self, elem: T) {
+        self.push_back(elem);
+    }
 }
 
 impl<I: IntoIterator> SpecExtend<I> for LinkedList<I::Item> {
@@ -1767,6 +1772,11 @@ impl<'a, T: 'a + Copy> Extend<&'a T> for LinkedList<T> {
     fn extend<I: IntoIterator<Item = &'a T>>(&mut self, iter: I) {
         self.extend(iter.into_iter().cloned());
     }
+
+    #[inline]
+    fn extend_one(&mut self, &elem: &'a T) {
+        self.push_back(elem);
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]

src/liballoc/collections/vec_deque.rs

@@ -2881,13 +2881,33 @@ impl<A> Extend<A> for VecDeque<A> {
             }
         }
     }
+
+    #[inline]
+    fn extend_one(&mut self, elem: A) {
+        self.push_back(elem);
+    }
+
+    #[inline]
+    fn extend_reserve(&mut self, additional: usize) {
+        self.reserve(additional);
+    }
 }
 
 #[stable(feature = "extend_ref", since = "1.2.0")]
 impl<'a, T: 'a + Copy> Extend<&'a T> for VecDeque<T> {
     fn extend<I: IntoIterator<Item = &'a T>>(&mut self, iter: I) {
         self.extend(iter.into_iter().cloned());
     }
+
+    #[inline]
+    fn extend_one(&mut self, &elem: &T) {
+        self.push_back(elem);
+    }
+
+    #[inline]
+    fn extend_reserve(&mut self, additional: usize) {
+        self.reserve(additional);
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]

src/liballoc/lib.rs

@@ -93,6 +93,7 @@
 #![feature(container_error_extra)]
 #![feature(dropck_eyepatch)]
 #![feature(exact_size_is_empty)]
+#![feature(extend_one)]
 #![feature(fmt_internals)]
 #![feature(fn_traits)]
 #![feature(fundamental)]

src/liballoc/string.rs

@@ -1799,34 +1799,69 @@ impl Extend<char> for String {
         self.reserve(lower_bound);
         iterator.for_each(move |c| self.push(c));
     }
+
+    #[inline]
+    fn extend_one(&mut self, c: char) {
+        self.push(c);
+    }
+
+    #[inline]
+    fn extend_reserve(&mut self, additional: usize) {
+        self.reserve(additional);
+    }
 }
 
 #[stable(feature = "extend_ref", since = "1.2.0")]
 impl<'a> Extend<&'a char> for String {
     fn extend<I: IntoIterator<Item = &'a char>>(&mut self, iter: I) {
         self.extend(iter.into_iter().cloned());
     }
+
+    #[inline]
+    fn extend_one(&mut self, &c: &'a char) {
+        self.push(c);
+    }
+
+    #[inline]
+    fn extend_reserve(&mut self, additional: usize) {
+        self.reserve(additional);
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<'a> Extend<&'a str> for String {
     fn extend<I: IntoIterator<Item = &'a str>>(&mut self, iter: I) {
         iter.into_iter().for_each(move |s| self.push_str(s));
     }
+
+    #[inline]
+    fn extend_one(&mut self, s: &'a str) {
+        self.push_str(s);
+    }
 }
 
 #[stable(feature = "extend_string", since = "1.4.0")]
 impl Extend<String> for String {
     fn extend<I: IntoIterator<Item = String>>(&mut self, iter: I) {
         iter.into_iter().for_each(move |s| self.push_str(&s));
     }
+
+    #[inline]
+    fn extend_one(&mut self, s: String) {
+        self.push_str(&s);
+    }
 }
 
 #[stable(feature = "herd_cows", since = "1.19.0")]
 impl<'a> Extend<Cow<'a, str>> for String {
     fn extend<I: IntoIterator<Item = Cow<'a, str>>>(&mut self, iter: I) {
         iter.into_iter().for_each(move |s| self.push_str(&s));
     }
+
+    #[inline]
+    fn extend_one(&mut self, s: Cow<'a, str>) {
+        self.push_str(&s);
+    }
 }
 
 /// A convenience impl that delegates to the impl for `&str`.

src/liballoc/vec.rs

@@ -2045,6 +2045,16 @@ 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 extend_one(&mut self, item: T) {
+        self.push(item);
+    }
+
+    #[inline]
+    fn extend_reserve(&mut self, additional: usize) {
+        self.reserve(additional);
+    }
 }
 
 // Specialization trait used for Vec::from_iter and Vec::extend
@@ -2316,6 +2326,16 @@ impl<'a, T: 'a + Copy> Extend<&'a T> for Vec<T> {
     fn extend<I: IntoIterator<Item = &'a T>>(&mut self, iter: I) {
         self.spec_extend(iter.into_iter())
     }
+
+    #[inline]
+    fn extend_one(&mut self, &item: &'a T) {
+        self.push(item);
+    }
+
+    #[inline]
+    fn extend_reserve(&mut self, additional: usize) {
+        self.reserve(additional);
+    }
 }
 
 macro_rules! __impl_slice_eq1 {

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,11 +341,26 @@ 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")]
 impl Extend<()> for () {
     fn extend<T: IntoIterator<Item = ()>>(&mut self, iter: T) {
         iter.into_iter().for_each(drop)
     }
+    fn extend_one(&mut self, _item: ()) {}
 }

src/libcore/iter/traits/iterator.rs

@@ -1700,9 +1700,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);
                 }
             }
         }
@@ -2675,14 +2675,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)

src/librustc_data_structures/lib.rs

@@ -22,6 +22,7 @@
 #![feature(test)]
 #![feature(associated_type_bounds)]
 #![feature(thread_id_value)]
+#![feature(extend_one)]
 #![allow(rustc::default_hash_types)]
 
 #[macro_use]

src/librustc_data_structures/thin_vec.rs

@@ -53,6 +53,20 @@ impl<T> Extend<T> for ThinVec<T> {
             ThinVec(None) => *self = iter.into_iter().collect::<Vec<_>>().into(),
         }
     }
+
+    fn extend_one(&mut self, item: T) {
+        match *self {
+            ThinVec(Some(ref mut vec)) => vec.push(item),
+            ThinVec(None) => *self = vec![item].into(),
+        }
+    }
+
+    fn extend_reserve(&mut self, additional: usize) {
+        match *self {
+            ThinVec(Some(ref mut vec)) => vec.reserve(additional),
+            ThinVec(None) => *self = Vec::with_capacity(additional).into(),
+        }
+    }
 }
 
 impl<T: HashStable<CTX>, CTX> HashStable<CTX> for ThinVec<T> {

src/librustc_index/lib.rs

@@ -2,6 +2,7 @@
 #![feature(const_if_match)]
 #![feature(const_fn)]
 #![feature(const_panic)]
+#![feature(extend_one)]
 #![feature(unboxed_closures)]
 #![feature(test)]
 #![feature(fn_traits)]

src/librustc_index/vec.rs

@@ -736,6 +736,16 @@ impl<I: Idx, T> Extend<T> for IndexVec<I, T> {
     fn extend<J: IntoIterator<Item = T>>(&mut self, iter: J) {
         self.raw.extend(iter);
     }
+
+    #[inline]
+    fn extend_one(&mut self, item: T) {
+        self.raw.push(item);
+    }
+
+    #[inline]
+    fn extend_reserve(&mut self, additional: usize) {
+        self.raw.reserve(additional);
+    }
 }
 
 impl<I: Idx, T> FromIterator<T> for IndexVec<I, T> {

src/librustc_infer/lib.rs

@@ -16,6 +16,7 @@
 #![feature(bool_to_option)]
 #![feature(box_patterns)]
 #![feature(box_syntax)]
+#![feature(extend_one)]
 #![feature(never_type)]
 #![feature(or_patterns)]
 #![feature(range_is_empty)]

src/librustc_infer/traits/util.rs

@@ -81,6 +81,14 @@ impl Extend<ty::Predicate<'tcx>> for PredicateSet<'tcx> {
             self.insert(pred);
         }
     }
+
+    fn extend_one(&mut self, pred: T) {
+        self.insert(pred.as_ref());
+    }
+
+    fn extend_reserve(&mut self, additional: usize) {
+        Extend::<ty::Predicate<'tcx>>::extend_reserve(&mut self.set, additional);
+    }
 }
 
 ///////////////////////////////////////////////////////////////////////////