From 3ca6bb0b44c3e65dab07e12aec9efb277dc206f9 Mon Sep 17 00:00:00 2001 From: The 8472 Date: Sat, 15 Apr 2023 12:59:16 +0200 Subject: Expand in-place iteration specialization to Flatten, FlatMap and ArrayChunks --- library/alloc/src/collections/binary_heap/mod.rs | 11 +++- library/alloc/src/lib.rs | 1 + library/alloc/src/vec/in_place_collect.rs | 83 +++++++++++++++++++----- library/alloc/src/vec/into_iter.rs | 12 +++- 4 files changed, 86 insertions(+), 21 deletions(-) (limited to 'library/alloc/src') diff --git a/library/alloc/src/collections/binary_heap/mod.rs b/library/alloc/src/collections/binary_heap/mod.rs index 66573b90db9..ad2c4e483a3 100644 --- a/library/alloc/src/collections/binary_heap/mod.rs +++ b/library/alloc/src/collections/binary_heap/mod.rs @@ -145,7 +145,7 @@ use core::alloc::Allocator; use core::fmt; -use core::iter::{FusedIterator, InPlaceIterable, SourceIter, TrustedLen}; +use core::iter::{FusedIterator, InPlaceIterable, SourceIter, TrustedFused, TrustedLen}; use core::mem::{self, swap, ManuallyDrop}; use core::num::NonZeroUsize; use core::ops::{Deref, DerefMut}; @@ -1540,6 +1540,10 @@ impl ExactSizeIterator for IntoIter { #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for IntoIter {} +#[doc(hidden)] +#[unstable(issue = "none", feature = "trusted_fused")] +unsafe impl TrustedFused for IntoIter {} + #[stable(feature = "default_iters", since = "1.70.0")] impl Default for IntoIter { /// Creates an empty `binary_heap::IntoIter`. @@ -1569,7 +1573,10 @@ unsafe impl SourceIter for IntoIter { #[unstable(issue = "none", feature = "inplace_iteration")] #[doc(hidden)] -unsafe impl InPlaceIterable for IntoIter {} +unsafe impl InPlaceIterable for IntoIter { + const EXPAND_BY: Option = NonZeroUsize::new(1); + const MERGE_BY: Option = NonZeroUsize::new(1); +} unsafe impl AsVecIntoIter for IntoIter { type Item = I; diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs index e43b6ac4039..c9fee435f1e 100644 --- a/library/alloc/src/lib.rs +++ b/library/alloc/src/lib.rs @@ -155,6 +155,7 @@ #![feature(std_internals)] #![feature(str_internals)] #![feature(strict_provenance)] +#![feature(trusted_fused)] #![feature(trusted_len)] #![feature(trusted_random_access)] #![feature(try_trait_v2)] diff --git a/library/alloc/src/vec/in_place_collect.rs b/library/alloc/src/vec/in_place_collect.rs index 5ecd0479971..d174074902a 100644 --- a/library/alloc/src/vec/in_place_collect.rs +++ b/library/alloc/src/vec/in_place_collect.rs @@ -137,44 +137,73 @@ //! } //! vec.truncate(write_idx); //! ``` +use crate::alloc::{handle_alloc_error, Global}; +use core::alloc::Allocator; +use core::alloc::Layout; use core::iter::{InPlaceIterable, SourceIter, TrustedRandomAccessNoCoerce}; use core::mem::{self, ManuallyDrop, SizedTypeProperties}; -use core::ptr::{self}; +use core::num::NonZeroUsize; +use core::ptr::{self, NonNull}; use super::{InPlaceDrop, InPlaceDstBufDrop, SpecFromIter, SpecFromIterNested, Vec}; -/// Specialization marker for collecting an iterator pipeline into a Vec while reusing the -/// source allocation, i.e. executing the pipeline in place. -#[rustc_unsafe_specialization_marker] -pub(super) trait InPlaceIterableMarker {} +const fn in_place_collectible( + step_merge: Option, + step_expand: Option, +) -> bool { + if DEST::IS_ZST || mem::align_of::() != mem::align_of::() { + return false; + } + + match (step_merge, step_expand) { + (Some(step_merge), Some(step_expand)) => { + // At least N merged source items -> at most M expanded destination items + // e.g. + // - 1 x [u8; 4] -> 4x u8, via flatten + // - 4 x u8 -> 1x [u8; 4], via array_chunks + mem::size_of::() * step_merge.get() == mem::size_of::() * step_expand.get() + } + // Fall back to other from_iter impls if an overflow occured in the step merge/expansion + // tracking. + _ => false, + } +} -impl InPlaceIterableMarker for T where T: InPlaceIterable {} +/// This provides a shorthand for the source type since local type aliases aren't a thing. +#[rustc_specialization_trait] +trait InPlaceCollect: SourceIter + InPlaceIterable { + type Src; +} + +impl InPlaceCollect for T +where + T: SourceIter + InPlaceIterable, +{ + type Src = <::Source as AsVecIntoIter>::Item; +} impl SpecFromIter for Vec where - I: Iterator + SourceIter + InPlaceIterableMarker, + I: Iterator + InPlaceCollect, + ::Source: AsVecIntoIter, { default fn from_iter(mut iterator: I) -> Self { // See "Layout constraints" section in the module documentation. We rely on const // optimization here since these conditions currently cannot be expressed as trait bounds - if T::IS_ZST - || mem::size_of::() - != mem::size_of::<<::Source as AsVecIntoIter>::Item>() - || mem::align_of::() - != mem::align_of::<<::Source as AsVecIntoIter>::Item>() - { + if const { !in_place_collectible::(I::MERGE_BY, I::EXPAND_BY) } { // fallback to more generic implementations return SpecFromIterNested::from_iter(iterator); } - let (src_buf, src_ptr, dst_buf, dst_end, cap) = unsafe { + let (src_buf, src_ptr, src_cap, mut dst_buf, dst_end, dst_cap) = unsafe { let inner = iterator.as_inner().as_into_iter(); ( inner.buf.as_ptr(), inner.ptr, + inner.cap, inner.buf.as_ptr() as *mut T, inner.end as *const T, - inner.cap, + inner.cap * mem::size_of::() / mem::size_of::(), ) }; @@ -203,11 +232,31 @@ where // Note: This access to the source wouldn't be allowed by the TrustedRandomIteratorNoCoerce // contract (used by SpecInPlaceCollect below). But see the "O(1) collect" section in the // module documentation why this is ok anyway. - let dst_guard = InPlaceDstBufDrop { ptr: dst_buf, len, cap }; + let dst_guard = InPlaceDstBufDrop { ptr: dst_buf, len, cap: dst_cap }; src.forget_allocation_drop_remaining(); mem::forget(dst_guard); - let vec = unsafe { Vec::from_raw_parts(dst_buf, len, cap) }; + // Adjust the allocation size if the source had a capacity in bytes that wasn't a multiple + // of the destination type size. + // Since the discrepancy should generally be small this should only result in some + // bookkeeping updates and no memmove. + if const { mem::size_of::() > mem::size_of::() } + && src_cap * mem::size_of::() != dst_cap * mem::size_of::() + { + let alloc = Global; + unsafe { + let new_layout = Layout::array::(dst_cap).unwrap(); + let result = alloc.shrink( + NonNull::new_unchecked(dst_buf as *mut u8), + Layout::array::(src_cap).unwrap(), + new_layout, + ); + let Ok(reallocated) = result else { handle_alloc_error(new_layout) }; + dst_buf = reallocated.as_ptr() as *mut T; + } + } + + let vec = unsafe { Vec::from_raw_parts(dst_buf, len, dst_cap) }; vec } diff --git a/library/alloc/src/vec/into_iter.rs b/library/alloc/src/vec/into_iter.rs index b2db2fdfd18..40b93b34e18 100644 --- a/library/alloc/src/vec/into_iter.rs +++ b/library/alloc/src/vec/into_iter.rs @@ -7,7 +7,8 @@ use crate::raw_vec::RawVec; use core::array; use core::fmt; use core::iter::{ - FusedIterator, InPlaceIterable, SourceIter, TrustedLen, TrustedRandomAccessNoCoerce, + FusedIterator, InPlaceIterable, SourceIter, TrustedFused, TrustedLen, + TrustedRandomAccessNoCoerce, }; use core::marker::PhantomData; use core::mem::{self, ManuallyDrop, MaybeUninit, SizedTypeProperties}; @@ -339,6 +340,10 @@ impl ExactSizeIterator for IntoIter { #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for IntoIter {} +#[doc(hidden)] +#[unstable(issue = "none", feature = "trusted_fused")] +unsafe impl TrustedFused for IntoIter {} + #[unstable(feature = "trusted_len", issue = "37572")] unsafe impl TrustedLen for IntoIter {} @@ -423,7 +428,10 @@ unsafe impl<#[may_dangle] T, A: Allocator> Drop for IntoIter { // also refer to the vec::in_place_collect module documentation to get an overview #[unstable(issue = "none", feature = "inplace_iteration")] #[doc(hidden)] -unsafe impl InPlaceIterable for IntoIter {} +unsafe impl InPlaceIterable for IntoIter { + const EXPAND_BY: Option = NonZeroUsize::new(1); + const MERGE_BY: Option = NonZeroUsize::new(1); +} #[unstable(issue = "none", feature = "inplace_iteration")] #[doc(hidden)] -- cgit 1.4.1-3-g733a5