From 54556f49d368ba96a92fd6ec352cf6b12e9542ef Mon Sep 17 00:00:00 2001 From: Chayim Refael Friedman Date: Sun, 7 Jul 2024 06:58:52 +0300 Subject: 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. --- library/alloc/src/vec/mod.rs | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) (limited to 'library/alloc/src/vec') 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 Extend for Vec { 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 Vec { @@ -3244,6 +3254,16 @@ impl<'a, T: Copy + 'a, A: Allocator> Extend<&'a T> for Vec { 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). -- cgit 1.4.1-3-g733a5