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 an unsafe method on `Extend` that doesn't check for growth (`extend_one_unchecked()`). I've marked it as perma-unstable currently, although we may want to expose it in the future so collections outside of std can benefit from it. Then specialize `Extend for (A, B)` for `TrustedLen` to call it.

It may seem that an alternative way of implementing this is to have a semi-public trait (`#[doc(hidden)]` public, so collections outside of core can implement it) for `extend()` inside tuples, and specialize it from collections. However, it is impossible due to limitations of `min_specialization`.

A concern that may arise with the current approach 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.

4 files changed, 53 insertions, 11 deletions

diff --git a/library/alloc/src/collections/vec_deque/mod.rs b/library/alloc/src/collections/vec_deque/mod.rs
index 61d05afccfd..a07f250d7d8 100644
--- a/library/alloc/src/collections/vec_deque/mod.rs
+++ b/library/alloc/src/collections/vec_deque/mod.rs
@@ -164,6 +164,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 {
@@ -2911,6 +2925,14 @@ impl<T, A: Allocator> Extend<T> for VecDeque<T, A> {
     fn extend_reserve(&mut self, additional: usize) {
         self.reserve(additional);
     }
+
+    #[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")]
@@ -2928,6 +2950,14 @@ impl<'a, T: 'a + Copy, A: Allocator> Extend<&'a T> for VecDeque<T, A> {
     fn extend_reserve(&mut self, additional: usize) {
         self.reserve(additional);
     }
+
+    #[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")]
diff --git a/library/alloc/src/collections/vec_deque/spec_extend.rs b/library/alloc/src/collections/vec_deque/spec_extend.rs
index dccf40ccb38..6a89abc3ef9 100644
--- a/library/alloc/src/collections/vec_deque/spec_extend.rs
+++ b/library/alloc/src/collections/vec_deque/spec_extend.rs
@@ -21,21 +21,12 @@ 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() {
@@ -43,7 +34,7 @@ where
                     return;
                 };
                 // SAFETY: The loop condition guarantees that `self.len() < self.capacity()`.
-                unsafe { push_unchecked(self, element) };
+                unsafe { self.push_unchecked(element) };
             }
         }
     }
diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs
index 703538d0f35..3237d8654c2 100644
--- a/library/alloc/src/lib.rs
+++ b/library/alloc/src/lib.rs
@@ -124,6 +124,7 @@
 #![feature(error_generic_member_access)]
 #![feature(exact_size_is_empty)]
 #![feature(extend_one)]
+#![feature(extend_one_unchecked)]
 #![feature(fmt_internals)]
 #![feature(fn_traits)]
 #![feature(hasher_prefixfree_extras)]
diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs
index f1706e31bb8..6e9b017ad75 100644
--- a/library/alloc/src/vec/mod.rs
+++ b/library/alloc/src/vec/mod.rs
@@ -3048,6 +3048,16 @@ impl<T, A: Allocator> Extend<T> for Vec<T, A> {
     fn extend_reserve(&mut self, additional: usize) {
         self.reserve(additional);
     }
+
+    #[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> {
@@ -3244,6 +3254,16 @@ impl<'a, T: Copy + 'a, A: Allocator> Extend<&'a T> for Vec<T, A> {
     fn extend_reserve(&mut self, additional: usize) {
         self.reserve(additional);
     }
+
+    #[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).