From 54556f49d368ba96a92fd6ec352cf6b12e9542ef Mon Sep 17 00:00:00 2001 From: Chayim Refael Friedman Date: Sun, 7 Jul 2024 06:58:52 +0300 Subject: [PATCH] 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. --- .../alloc/src/collections/vec_deque/mod.rs | 30 +++++ .../src/collections/vec_deque/spec_extend.rs | 13 +- library/alloc/src/lib.rs | 1 + library/alloc/src/vec/mod.rs | 20 +++ library/core/src/iter/traits/collect.rs | 122 +++++++++++++++--- 5 files changed, 160 insertions(+), 26 deletions(-) diff --git a/library/alloc/src/collections/vec_deque/mod.rs b/library/alloc/src/collections/vec_deque/mod.rs index 61d05afccfd1a..a07f250d7d88c 100644 --- a/library/alloc/src/collections/vec_deque/mod.rs +++ b/library/alloc/src/collections/vec_deque/mod.rs @@ -164,6 +164,20 @@ impl VecDeque { 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 Extend for VecDeque { 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 { 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 dccf40ccb38aa..6a89abc3ef9b6 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(deque: &mut VecDeque, 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 703538d0f355d..3237d8654c281 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 f1706e31bb80a..6e9b017ad75cf 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). diff --git a/library/core/src/iter/traits/collect.rs b/library/core/src/iter/traits/collect.rs index 61a457890134c..86660f2e375c3 100644 --- a/library/core/src/iter/traits/collect.rs +++ b/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 @@ -460,6 +462,27 @@ pub trait Extend { fn extend_reserve(&mut self, additional: usize) { let _ = additional; } + + /// Extends a collection with one element, without checking there is enough capacity for it. + /// + /// # Safety + /// + /// **For callers:** 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`. + /// + /// **For implementors:** 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. + + // This method is for internal usage only. It is only on the trait because of specialization's limitations. + #[unstable(feature = "extend_one_unchecked", issue = "none")] + #[doc(hidden)] + unsafe fn extend_one_unchecked(&mut self, item: A) + where + Self: Sized, + { + self.extend_one(item); + } } #[stable(feature = "extend_for_unit", since = "1.28.0")] @@ -499,33 +522,102 @@ where fn extend>(&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); + } + + 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( + iter: impl Iterator, + a: &mut ExtendA, + b: &mut ExtendB, +) where + ExtendA: Extend, + ExtendB: Extend, +{ + fn extend<'a, A, B>( + a: &'a mut impl Extend, + b: &'a mut impl Extend, + ) -> 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 { + fn extend(self, a: &mut A, b: &mut B); +} +impl SpecTupleExtend for Iter +where + ExtendA: Extend, + ExtendB: Extend, + Iter: Iterator, +{ + default fn extend(self, a: &mut ExtendA, b: &mut ExtendB) { + default_extend_tuple(self, a, b); + } +} + +impl SpecTupleExtend for Iter +where + ExtendA: Extend, + ExtendB: Extend, + Iter: TrustedLen, +{ + fn extend(self, a: &mut ExtendA, b: &mut ExtendB) { fn extend<'a, A, B>( a: &'a mut impl Extend, b: &'a mut impl Extend, ) -> 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)); } }