mikros/rust
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Specialize TrustedLen for Iterator::unzip()

Browse Source 
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.
This commit is contained in:
Chayim Refael Friedman 2024-07-07 06:58:52 +03:00
parent 6ba80a9f8d
commit 54556f49d3
5 changed files with 163 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
library/alloc/src/collections/vec_deque/mod.rs
View File

@ -164,6 +164,20 @@ fn ptr(&self) -> *mut T {
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 @@ fn extend_one(&mut self, elem: T) {
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 @@ fn extend_one(&mut self, &elem: &'a T) {
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")]

13
library/alloc/src/collections/vec_deque/spec_extend.rs
View File

@ -21,21 +21,12 @@ impl<T, I, A: Allocator> SpecExtend<T, I> for VecDeque<T, A>
// 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 @@ unsafe fn push_unchecked<T, A: Allocator>(deque: &mut VecDeque<T, A>, element: T
return;
};
// SAFETY: The loop condition guarantees that `self.len() < self.capacity()`.
unsafe { push_unchecked(self, element) };
unsafe { self.push_unchecked(element) };
}
}
}

1
library/alloc/src/lib.rs
View File

@ -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)]

20
library/alloc/src/vec/mod.rs
View File

@ -3048,6 +3048,16 @@ fn extend_one(&mut self, item: T) {
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 @@ fn extend_one(&mut self, &item: &'a T) {
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).

104
library/core/src/iter/traits/collect.rs
View File

@ -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 @@ fn extend_one(&mut self, item: A) {
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,7 +522,36 @@ impl<A, B, ExtendA, ExtendB> Extend<(A, B)> for (ExtendA, ExtendB)
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);
}
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>,
@ -519,13 +571,53 @@ fn extend<'a, A, B>(
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);
trait SpecTupleExtend<A, B> {
fn extend(self, a: &mut A, b: &mut B);
}
fn extend_reserve(&mut self, additional: usize) {
self.0.extend_reserve(additional);
self.1.extend_reserve(additional);
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 {
// 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, 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);
}
self.fold((), extend(a, b));
}
}