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.

1 files changed, 20 insertions, 0 deletions

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).