Specialize TrustedLen for Iterator::unzip()

Don't check the capacity every time (and also for `Extend` for tuples, as this is how `unzip()` is implemented).

I did this with a semi-public (`#[doc(hidden)]` public, so collections outside of core can implement it) specialization trait that has a method (`extend_one_unchecked()`) that doesn't check for growth. Then specialize `Extend for (A, B)` for `TrustedLen` to call it, and specialize it for collections (only `Vec` and `VecDequq` from the collection in the standard library could benefit from this)

An alternative way of implementing this is to have this method on `Extend` directly, but this was rejected by @the8472 in a review (rust-lang#123253 (comment)).

I didn't mark the trait as unsafe; a concern that may arise from that is that implementing `extend_one_unchecked()` correctly must also incur implementing `extend_reserve()`, otherwise you can have UB. This is a somewhat non-local safety invariant. However, I believe this is fine, since to have actual UB you must have unsafe code inside your `extend_one_unchecked()` that makes incorrect assumption, *and* not implement `extend_reserve()`. I've also documented this requirement.

Author: ChayimFriedman2

library/alloc/src/collections/vec_deque/mod.rs

@@ -10,7 +10,7 @@
 use core::cmp::{self, Ordering};
 use core::fmt;
 use core::hash::{Hash, Hasher};
-use core::iter::{repeat_n, repeat_with, ByRefSized};
+use core::iter::{repeat_n, repeat_with, ByRefSized, ExtendUnchecked};
 use core::mem::{ManuallyDrop, SizedTypeProperties};
 use core::ops::{Index, IndexMut, Range, RangeBounds};
 use core::ptr;
@@ -160,6 +160,20 @@ impl<T, A: Allocator> VecDeque<T, A> {
         self.buf.ptr()
     }
 
+    /// Appends an element to the buffer.
+    ///
+    /// # Safety
+    ///
+    /// May only be called if `deque.len() < deque.capacity()`
+    #[inline]
+    unsafe fn push_unchecked(&mut self, element: T) {
+        // SAFETY: Because of the precondition, it's guaranteed that there is space
+        // in the logical array after the last element.
+        unsafe { self.buffer_write(self.to_physical_idx(self.len), element) };
+        // This can't overflow because `deque.len() < deque.capacity() <= usize::MAX`.
+        self.len += 1;
+    }
+
     /// Moves an element out of the buffer
     #[inline]
     unsafe fn buffer_read(&mut self, off: usize) -> T {
@@ -2837,6 +2851,16 @@ impl<T, A: Allocator> Extend<T> for VecDeque<T, A> {
     }
 }
 
+impl<T, A: Allocator> ExtendUnchecked<T> for VecDeque<T, A> {
+    #[inline]
+    unsafe fn extend_one_unchecked(&mut self, item: T) {
+        // SAFETY: Our preconditions ensure the space has been reserved, and `extend_reserve` is implemented correctly.
+        unsafe {
+            self.push_unchecked(item);
+        }
+    }
+}
+
 #[stable(feature = "extend_ref", since = "1.2.0")]
 impl<'a, T: 'a + Copy, A: Allocator> Extend<&'a T> for VecDeque<T, A> {
     fn extend<I: IntoIterator<Item = &'a T>>(&mut self, iter: I) {
@@ -2854,6 +2878,16 @@ impl<'a, T: 'a + Copy, A: Allocator> Extend<&'a T> for VecDeque<T, A> {
     }
 }
 
+impl<'a, T: 'a + Copy, A: Allocator> ExtendUnchecked<&'a T> for VecDeque<T, A> {
+    #[inline]
+    unsafe fn extend_one_unchecked(&mut self, &item: &'a T) {
+        // SAFETY: Our preconditions ensure the space has been reserved, and `extend_reserve` is implemented correctly.
+        unsafe {
+            self.push_unchecked(item);
+        }
+    }
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<T: fmt::Debug, A: Allocator> fmt::Debug for VecDeque<T, A> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {

library/alloc/src/collections/vec_deque/spec_extend.rs

@@ -21,29 +21,20 @@ where
         //     self.push_back(item);
         // }
 
-        // May only be called if `deque.len() < deque.capacity()`
-        unsafe fn push_unchecked<T, A: Allocator>(deque: &mut VecDeque<T, A>, element: T) {
-            // SAFETY: Because of the precondition, it's guaranteed that there is space
-            // in the logical array after the last element.
-            unsafe { deque.buffer_write(deque.to_physical_idx(deque.len), element) };
-            // This can't overflow because `deque.len() < deque.capacity() <= usize::MAX`.
-            deque.len += 1;
-        }
-
         while let Some(element) = iter.next() {
             let (lower, _) = iter.size_hint();
             self.reserve(lower.saturating_add(1));
 
             // SAFETY: We just reserved space for at least one element.
-            unsafe { push_unchecked(self, element) };
+            unsafe { self.push_unchecked(element) };
 
             // Inner loop to avoid repeatedly calling `reserve`.
             while self.len < self.capacity() {
                 let Some(element) = iter.next() else {
                     return;
                 };
                 // SAFETY: The loop condition guarantees that `self.len() < self.capacity()`.
-                unsafe { push_unchecked(self, element) };
+                unsafe { self.push_unchecked(element) };
             }
         }
     }

library/alloc/src/lib.rs

@@ -128,6 +128,7 @@
 #![feature(error_in_core)]
 #![feature(exact_size_is_empty)]
 #![feature(extend_one)]
+#![feature(extend_unchecked)]
 #![feature(fmt_internals)]
 #![feature(fn_traits)]
 #![feature(generic_nonzero)]

library/alloc/src/vec/mod.rs

@@ -59,7 +59,7 @@ use core::cmp::Ordering;
 use core::fmt;
 use core::hash::{Hash, Hasher};
 #[cfg(not(no_global_oom_handling))]
-use core::iter;
+use core::iter::{self, ExtendUnchecked};
 use core::marker::PhantomData;
 use core::mem::{self, ManuallyDrop, MaybeUninit, SizedTypeProperties};
 use core::ops::{self, Index, IndexMut, Range, RangeBounds};
@@ -3000,6 +3000,19 @@ impl<T, A: Allocator> Extend<T> for Vec<T, A> {
     }
 }
 
+#[cfg(not(no_global_oom_handling))]
+impl<T, A: Allocator> ExtendUnchecked<T> for Vec<T, A> {
+    #[inline]
+    unsafe fn extend_one_unchecked(&mut self, item: T) {
+        // SAFETY: Our preconditions ensure the space has been reserved, and `extend_reserve` is implemented correctly.
+        unsafe {
+            let len = self.len();
+            ptr::write(self.as_mut_ptr().add(len), item);
+            self.set_len(len + 1);
+        }
+    }
+}
+
 impl<T, A: Allocator> Vec<T, A> {
     // leaf method to which various SpecFrom/SpecExtend implementations delegate when
     // they have no further optimizations to apply
@@ -3196,6 +3209,19 @@ impl<'a, T: Copy + 'a, A: Allocator> Extend<&'a T> for Vec<T, A> {
     }
 }
 
+#[cfg(not(no_global_oom_handling))]
+impl<'a, T: Copy + 'a, A: Allocator> ExtendUnchecked<&'a T> for Vec<T, A> {
+    #[inline]
+    unsafe fn extend_one_unchecked(&mut self, &item: &'a T) {
+        // SAFETY: Our preconditions ensure the space has been reserved, and `extend_reserve` is implemented correctly.
+        unsafe {
+            let len = self.len();
+            ptr::write(self.as_mut_ptr().add(len), item);
+            self.set_len(len + 1);
+        }
+    }
+}
+
 /// Implements comparison of vectors, [lexicographically](Ord#lexicographical-comparison).
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<T, A1, A2> PartialOrd<Vec<T, A2>> for Vec<T, A1>

library/core/src/iter/mod.rs

@@ -413,6 +413,9 @@ pub use self::sources::{repeat_with, RepeatWith};
 #[stable(feature = "iter_successors", since = "1.34.0")]
 pub use self::sources::{successors, Successors};
 
+#[doc(hidden)]
+#[unstable(feature = "extend_unchecked", issue = "none")]
+pub use self::traits::ExtendUnchecked;
 #[stable(feature = "fused", since = "1.26.0")]
 pub use self::traits::FusedIterator;
 #[unstable(issue = "none", feature = "inplace_iteration")]

library/core/src/iter/traits/collect.rs

@@ -1,3 +1,5 @@
+use super::TrustedLen;
+
 /// Conversion from an [`Iterator`].
 ///
 /// By implementing `FromIterator` for a type, you define how it will be
@@ -429,6 +431,30 @@ pub trait Extend<A> {
     }
 }
 
+// This trait is exported so collections outside of `core` can implement it.
+#[doc(hidden)]
+#[unstable(feature = "extend_unchecked", issue = "none")]
+pub trait ExtendUnchecked<A>: Extend<A> {
+    /// Extends a collection with one element, without checking there is enough capacity for it.
+    ///
+    /// **Note:** For a collection to unsafely rely on this method's safety precondition (that is,
+    /// invoke UB if they are violated), it must implement `extend_reserve` correctly. In other words,
+    /// callers may assume that if they `extend_reserve`ed enough space they can call this method.
+    ///
+    /// # Safety
+    ///
+    /// This must only be called when we know the collection has enough capacity to contain the new item,
+    /// for example because we previously called `extend_reserve`.
+    unsafe fn extend_one_unchecked(&mut self, item: A);
+}
+
+#[unstable(feature = "extend_unchecked", issue = "none")]
+impl<A, T: Extend<A>> ExtendUnchecked<A> for T {
+    default unsafe fn extend_one_unchecked(&mut self, item: A) {
+        self.extend_one(item);
+    }
+}
+
 #[stable(feature = "extend_for_unit", since = "1.28.0")]
 impl Extend<()> for () {
     fn extend<T: IntoIterator<Item = ()>>(&mut self, iter: T) {
@@ -466,33 +492,108 @@ where
     fn extend<T: IntoIterator<Item = (A, B)>>(&mut self, into_iter: T) {
         let (a, b) = self;
         let iter = into_iter.into_iter();
+        SpecTupleExtend::extend(iter, a, b);
+    }
+
+    fn extend_one(&mut self, item: (A, B)) {
+        self.0.extend_one(item.0);
+        self.1.extend_one(item.1);
+    }
 
+    fn extend_reserve(&mut self, additional: usize) {
+        self.0.extend_reserve(additional);
+        self.1.extend_reserve(additional);
+    }
+}
+
+impl<A, B, ExtendA, ExtendB> ExtendUnchecked<(A, B)> for (ExtendA, ExtendB)
+where
+    ExtendA: Extend<A>,
+    ExtendB: Extend<B>,
+{
+    unsafe fn extend_one_unchecked(&mut self, item: (A, B)) {
+        // SAFETY: Those are our safety preconditions, and we correctly forward `extend_reserve`.
+        unsafe {
+            self.0.extend_one_unchecked(item.0);
+            self.1.extend_one_unchecked(item.1);
+        }
+    }
+}
+
+fn default_extend_tuple<A, B, ExtendA, ExtendB>(
+    iter: impl Iterator<Item = (A, B)>,
+    a: &mut ExtendA,
+    b: &mut ExtendB,
+) where
+    ExtendA: Extend<A>,
+    ExtendB: Extend<B>,
+{
+    fn extend<'a, A, B>(
+        a: &'a mut impl Extend<A>,
+        b: &'a mut impl Extend<B>,
+    ) -> impl FnMut((), (A, B)) + 'a {
+        move |(), (t, u)| {
+            a.extend_one(t);
+            b.extend_one(u);
+        }
+    }
+
+    let (lower_bound, _) = iter.size_hint();
+    if lower_bound > 0 {
+        a.extend_reserve(lower_bound);
+        b.extend_reserve(lower_bound);
+    }
+
+    iter.fold((), extend(a, b));
+}
+
+trait SpecTupleExtend<A, B> {
+    fn extend(self, a: &mut A, b: &mut B);
+}
+
+impl<A, B, ExtendA, ExtendB, Iter> SpecTupleExtend<ExtendA, ExtendB> for Iter
+where
+    ExtendA: Extend<A>,
+    ExtendB: Extend<B>,
+    Iter: Iterator<Item = (A, B)>,
+{
+    default fn extend(self, a: &mut ExtendA, b: &mut ExtendB) {
+        default_extend_tuple(self, a, b);
+    }
+}
+
+impl<A, B, ExtendA, ExtendB, Iter> SpecTupleExtend<ExtendA, ExtendB> for Iter
+where
+    ExtendA: Extend<A>,
+    ExtendB: Extend<B>,
+    Iter: TrustedLen<Item = (A, B)>,
+{
+    fn extend(self, a: &mut ExtendA, b: &mut ExtendB) {
         fn extend<'a, A, B>(
             a: &'a mut impl Extend<A>,
             b: &'a mut impl Extend<B>,
         ) -> impl FnMut((), (A, B)) + 'a {
-            move |(), (t, u)| {
-                a.extend_one(t);
-                b.extend_one(u);
+            // SAFETY: We reserve enough space for the `size_hint`, and the iterator is `TrustedLen`
+            // so its `size_hint` is exact.
+            move |(), (t, u)| unsafe {
+                a.extend_one_unchecked(t);
+                b.extend_one_unchecked(u);
             }
         }
 
-        let (lower_bound, _) = iter.size_hint();
+        let (lower_bound, upper_bound) = self.size_hint();
+
+        if upper_bound.is_none() {
+            // We cannot reserve more than `usize::MAX` items, and this is likely to go out of memory anyway.
+            default_extend_tuple(self, a, b);
+            return;
+        }
+
         if lower_bound > 0 {
             a.extend_reserve(lower_bound);
             b.extend_reserve(lower_bound);
         }
 
-        iter.fold((), extend(a, b));
-    }
-
-    fn extend_one(&mut self, item: (A, B)) {
-        self.0.extend_one(item.0);
-        self.1.extend_one(item.1);
-    }
-
-    fn extend_reserve(&mut self, additional: usize) {
-        self.0.extend_reserve(additional);
-        self.1.extend_reserve(additional);
+        self.fold((), extend(a, b));
     }
 }

library/core/src/iter/traits/mod.rs

@@ -16,6 +16,9 @@ pub use self::{
     marker::{FusedIterator, TrustedLen},
 };
 
+#[doc(hidden)]
+#[unstable(feature = "extend_unchecked", issue = "none")]
+pub use self::collect::ExtendUnchecked;
 #[unstable(issue = "none", feature = "inplace_iteration")]
 pub use self::marker::InPlaceIterable;
 #[unstable(issue = "none", feature = "trusted_fused")]