diff --git a/library/alloc/src/collections/binary_heap/mod.rs b/library/alloc/src/collections/binary_heap/mod.rs index 66573b90db963..ad2c4e483a3bf 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 e43b6ac40391d..c9fee435f1e5d 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 5ecd0479971ea..e4f96fd764040 100644 --- a/library/alloc/src/vec/in_place_collect.rs +++ b/library/alloc/src/vec/in_place_collect.rs @@ -6,11 +6,11 @@ //! The specialization in this module applies to iterators in the shape of //! `source.adapter().adapter().adapter().collect::>()` //! where `source` is an owning iterator obtained from [`Vec`], [`Box<[T]>`][box] (by conversion to `Vec`) -//! or [`BinaryHeap`], the adapters each consume one or more items per step -//! (represented by [`InPlaceIterable`]), provide transitive access to `source` (via [`SourceIter`]) -//! and thus the underlying allocation. And finally the layouts of `T` and `U` must -//! have the same size and alignment, this is currently ensured via const eval instead of trait bounds -//! in the specialized [`SpecFromIter`] implementation. +//! or [`BinaryHeap`], the adapters guarantee to consume enough items per step to make room +//! for the results (represented by [`InPlaceIterable`]), provide transitive access to `source` +//! (via [`SourceIter`]) and thus the underlying allocation. +//! And finally there are alignment and size constriants to consider, this is currently ensured via +//! const eval instead of trait bounds in the specialized [`SpecFromIter`] implementation. //! //! [`BinaryHeap`]: crate::collections::BinaryHeap //! [box]: crate::boxed::Box @@ -35,11 +35,28 @@ //! the step of reading a value and getting a reference to write to. Instead raw pointers must be //! used on the reader and writer side. //! -//! That writes never clobber a yet-to-be-read item is ensured by the [`InPlaceIterable`] requirements. +//! That writes never clobber a yet-to-be-read items is ensured by the [`InPlaceIterable`] requirements. //! //! # Layout constraints //! -//! [`Allocator`] requires that `allocate()` and `deallocate()` have matching alignment and size. +//! When recycling an allocation between different types we must uphold the [`Allocator`] contract +//! which means that the input and output Layouts have to "fit". +//! +//! To complicate things further `InPlaceIterable` supports splitting or merging items into smaller/ +//! larger ones to enable (de)aggregation of arrays. +//! +//! Ultimately each step of the iterator must free up enough *bytes* in the source to make room +//! for the next output item. +//! If `T` and `U` have the same size no fixup is needed. +//! If `T`'s size is a multiple of `U`'s we can compensate by multiplying the capacity accordingly. +//! Otherwise the input capacity (and thus layout) in bytes may not be representable by the output +//! `Vec`. In that case `alloc.shrink()` is used to update the allocation's layout. +//! +//! Alignments of `T` must be the same or larger than `U`. Since alignments are always a power +//! of two _larger_ implies _is a multiple of_. +//! +//! See `in_place_collectible()` for the current conditions. +//! //! Additionally this specialization doesn't make sense for ZSTs as there is no reallocation to //! avoid and it would make pointer arithmetic more difficult. //! @@ -137,44 +154,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 occurred 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::(), ) }; @@ -196,18 +242,55 @@ where } // The ownership of the allocation and the new `T` values is temporarily moved into `dst_guard`. - // This is safe because `forget_allocation_drop_remaining` immediately forgets the allocation + // This is safe because + // * `forget_allocation_drop_remaining` immediately forgets the allocation // before any panic can occur in order to avoid any double free, and then proceeds to drop // any remaining values at the tail of the source. + // * the shrink either panics without invalidating the allocation, aborts or + // succeeds. In the last case we disarm the guard. // // 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(); + + // Adjust the allocation if the alignment didn't match or 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 { + let src_sz = mem::size_of::(); + src_sz > 0 && mem::size_of::() % src_sz != 0 + } && src_cap * mem::size_of::() != dst_cap * mem::size_of::()) + || const { mem::align_of::() != mem::align_of::() } + { + let alloc = Global; + unsafe { + // The old allocation exists, therefore it must have a valid layout. + let src_align = mem::align_of::(); + let src_size = mem::size_of::().unchecked_mul(src_cap); + let old_layout = Layout::from_size_align_unchecked(src_size, src_align); + + // The must be equal or smaller for in-place iteration to be possible + // therefore the new layout must be ≤ the old one and therefore valid. + let dst_align = mem::align_of::(); + let dst_size = mem::size_of::().unchecked_mul(dst_cap); + let new_layout = Layout::from_size_align_unchecked(dst_size, dst_align); + + let result = alloc.shrink( + NonNull::new_unchecked(dst_buf as *mut u8), + old_layout, + new_layout, + ); + let Ok(reallocated) = result else { handle_alloc_error(new_layout) }; + dst_buf = reallocated.as_ptr() as *mut T; + } + } + mem::forget(dst_guard); - let vec = unsafe { Vec::from_raw_parts(dst_buf, len, cap) }; + 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 b2db2fdfd18f1..40b93b34e18bb 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)] diff --git a/library/alloc/tests/lib.rs b/library/alloc/tests/lib.rs index aa7a331b368be..5bb531f0e6272 100644 --- a/library/alloc/tests/lib.rs +++ b/library/alloc/tests/lib.rs @@ -1,5 +1,6 @@ #![feature(allocator_api)] #![feature(alloc_layout_extra)] +#![feature(iter_array_chunks)] #![feature(assert_matches)] #![feature(btree_extract_if)] #![feature(cow_is_borrowed)] diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs index 9cb27899f1091..c545f47a5d15f 100644 --- a/library/alloc/tests/vec.rs +++ b/library/alloc/tests/vec.rs @@ -1,6 +1,7 @@ +use alloc::vec::Vec; use core::alloc::{Allocator, Layout}; -use core::assert_eq; -use core::iter::IntoIterator; +use core::{assert_eq, assert_ne}; +use core::iter::{IntoIterator, Iterator}; use core::num::NonZeroUsize; use core::ptr::NonNull; use std::alloc::System; @@ -1184,6 +1185,46 @@ fn test_from_iter_specialization_with_iterator_adapters() { assert_eq!(srcptr, sinkptr as *const usize); } +#[test] +fn test_in_place_specialization_step_up_down() { + fn assert_in_place_trait(_: &T) {} + let src = vec![[0u8; 4]; 256]; + let srcptr = src.as_ptr(); + let src_cap = src.capacity(); + let iter = src.into_iter().flatten(); + assert_in_place_trait(&iter); + let sink = iter.collect::>(); + let sinkptr = sink.as_ptr(); + assert_eq!(srcptr as *const u8, sinkptr); + assert_eq!(src_cap * 4, sink.capacity()); + + let iter = sink.into_iter().array_chunks::<4>(); + assert_in_place_trait(&iter); + let sink = iter.collect::>(); + let sinkptr = sink.as_ptr(); + assert_eq!(srcptr, sinkptr); + assert_eq!(src_cap, sink.capacity()); + + let mut src: Vec = Vec::with_capacity(17); + let src_bytes = src.capacity(); + src.resize(8, 0u8); + let sink: Vec<[u8; 4]> = src.into_iter().array_chunks::<4>().collect(); + let sink_bytes = sink.capacity() * 4; + assert_ne!(src_bytes, sink_bytes); + assert_eq!(sink.len(), 2); + + let src = vec![[0u8; 4]; 256]; + let srcptr = src.as_ptr(); + let iter = src + .into_iter() + .flat_map(|a| { + a.into_iter().map(|b| b.wrapping_add(1)) + }); + assert_in_place_trait(&iter); + let sink = iter.collect::>(); + assert_eq!(srcptr as *const u8, sink.as_ptr()); +} + #[test] fn test_from_iter_specialization_head_tail_drop() { let drop_count: Vec<_> = (0..=2).map(|_| Rc::new(())).collect(); diff --git a/library/core/src/iter/adapters/array_chunks.rs b/library/core/src/iter/adapters/array_chunks.rs index 13719c727e93f..319af4408d76d 100644 --- a/library/core/src/iter/adapters/array_chunks.rs +++ b/library/core/src/iter/adapters/array_chunks.rs @@ -1,5 +1,9 @@ use crate::array; -use crate::iter::{ByRefSized, FusedIterator, Iterator, TrustedRandomAccessNoCoerce}; +use crate::iter::adapters::SourceIter; +use crate::iter::{ + ByRefSized, FusedIterator, InPlaceIterable, Iterator, TrustedFused, TrustedRandomAccessNoCoerce, +}; +use crate::num::NonZeroUsize; use crate::ops::{ControlFlow, NeverShortCircuit, Try}; /// An iterator over `N` elements of the iterator at a time. @@ -159,6 +163,9 @@ where #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")] impl FusedIterator for ArrayChunks where I: FusedIterator {} +#[unstable(issue = "none", feature = "trusted_fused")] +unsafe impl TrustedFused for ArrayChunks where I: TrustedFused + Iterator {} + #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")] impl ExactSizeIterator for ArrayChunks where @@ -229,3 +236,28 @@ where accum } } + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl SourceIter for ArrayChunks +where + I: SourceIter + Iterator, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut I::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.iter) } + } +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl InPlaceIterable for ArrayChunks { + const EXPAND_BY: Option = I::EXPAND_BY; + const MERGE_BY: Option = const { + match (I::MERGE_BY, NonZeroUsize::new(N)) { + (Some(m), Some(n)) => m.checked_mul(n), + _ => None, + } + }; +} diff --git a/library/core/src/iter/adapters/enumerate.rs b/library/core/src/iter/adapters/enumerate.rs index 00c1c377bf9c2..92f465ccdb4e8 100644 --- a/library/core/src/iter/adapters/enumerate.rs +++ b/library/core/src/iter/adapters/enumerate.rs @@ -1,7 +1,7 @@ use crate::iter::adapters::{ zip::try_get_unchecked, SourceIter, TrustedRandomAccess, TrustedRandomAccessNoCoerce, }; -use crate::iter::{FusedIterator, InPlaceIterable, TrustedLen}; +use crate::iter::{FusedIterator, InPlaceIterable, TrustedFused, TrustedLen}; use crate::num::NonZeroUsize; use crate::ops::Try; @@ -243,6 +243,9 @@ where #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for Enumerate where I: FusedIterator {} +#[unstable(issue = "none", feature = "trusted_fused")] +unsafe impl TrustedFused for Enumerate {} + #[unstable(feature = "trusted_len", issue = "37572")] unsafe impl TrustedLen for Enumerate where I: TrustedLen {} @@ -261,7 +264,10 @@ where } #[unstable(issue = "none", feature = "inplace_iteration")] -unsafe impl InPlaceIterable for Enumerate {} +unsafe impl InPlaceIterable for Enumerate { + const EXPAND_BY: Option = I::EXPAND_BY; + const MERGE_BY: Option = I::MERGE_BY; +} #[stable(feature = "default_iters", since = "1.70.0")] impl Default for Enumerate { diff --git a/library/core/src/iter/adapters/filter.rs b/library/core/src/iter/adapters/filter.rs index 723657b9e43e4..882f3e3bc60f5 100644 --- a/library/core/src/iter/adapters/filter.rs +++ b/library/core/src/iter/adapters/filter.rs @@ -1,5 +1,6 @@ use crate::fmt; -use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable}; +use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedFused}; +use crate::num::NonZeroUsize; use crate::ops::Try; use core::array; use core::mem::{ManuallyDrop, MaybeUninit}; @@ -189,6 +190,9 @@ where #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for Filter where P: FnMut(&I::Item) -> bool {} +#[unstable(issue = "none", feature = "trusted_fused")] +unsafe impl TrustedFused for Filter {} + #[unstable(issue = "none", feature = "inplace_iteration")] unsafe impl SourceIter for Filter where @@ -204,4 +208,7 @@ where } #[unstable(issue = "none", feature = "inplace_iteration")] -unsafe impl InPlaceIterable for Filter where P: FnMut(&I::Item) -> bool {} +unsafe impl InPlaceIterable for Filter { + const EXPAND_BY: Option = I::EXPAND_BY; + const MERGE_BY: Option = I::MERGE_BY; +} diff --git a/library/core/src/iter/adapters/filter_map.rs b/library/core/src/iter/adapters/filter_map.rs index 693479977db51..c0bdc065b05dd 100644 --- a/library/core/src/iter/adapters/filter_map.rs +++ b/library/core/src/iter/adapters/filter_map.rs @@ -1,5 +1,6 @@ -use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable}; +use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedFused}; use crate::mem::{ManuallyDrop, MaybeUninit}; +use crate::num::NonZeroUsize; use crate::ops::{ControlFlow, Try}; use crate::{array, fmt}; @@ -188,6 +189,9 @@ where #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for FilterMap where F: FnMut(I::Item) -> Option {} +#[unstable(issue = "none", feature = "trusted_fused")] +unsafe impl TrustedFused for FilterMap {} + #[unstable(issue = "none", feature = "inplace_iteration")] unsafe impl SourceIter for FilterMap where @@ -203,7 +207,7 @@ where } #[unstable(issue = "none", feature = "inplace_iteration")] -unsafe impl InPlaceIterable for FilterMap where - F: FnMut(I::Item) -> Option -{ +unsafe impl InPlaceIterable for FilterMap { + const EXPAND_BY: Option = I::EXPAND_BY; + const MERGE_BY: Option = I::MERGE_BY; } diff --git a/library/core/src/iter/adapters/flatten.rs b/library/core/src/iter/adapters/flatten.rs index eee6e5bccecc7..09428350fd92a 100644 --- a/library/core/src/iter/adapters/flatten.rs +++ b/library/core/src/iter/adapters/flatten.rs @@ -1,7 +1,13 @@ -use crate::fmt; -use crate::iter::{DoubleEndedIterator, Fuse, FusedIterator, Iterator, Map, TrustedLen}; +use crate::iter::adapters::SourceIter; +use crate::iter::{ + Cloned, Copied, DoubleEndedIterator, Filter, FilterMap, Fuse, FusedIterator, InPlaceIterable, + Iterator, Map, TrustedFused, TrustedLen, +}; +use crate::iter::{Once, OnceWith}; use crate::num::NonZeroUsize; use crate::ops::{ControlFlow, Try}; +use crate::result; +use crate::{array, fmt, option}; /// An iterator that maps each element to an iterator, and yields the elements /// of the produced iterators. @@ -145,6 +151,91 @@ where { } +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl InPlaceIterable for FlatMap +where + I: InPlaceIterable, + U: BoundedSize + IntoIterator, +{ + const EXPAND_BY: Option = const { + match (I::EXPAND_BY, U::UPPER_BOUND) { + (Some(m), Some(n)) => m.checked_mul(n), + _ => None, + } + }; + const MERGE_BY: Option = I::MERGE_BY; +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl SourceIter for FlatMap +where + I: SourceIter + TrustedFused, + U: IntoIterator, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut I::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.inner.iter) } + } +} + +/// Marker trait for iterators/iterables which have a statically known upper +/// bound of the number of items they can produce. +/// +/// # Safety +/// +/// Implementations must not yield more elements than indicated by UPPER_BOUND if it is `Some`. +/// Used in specializations. Implementations must not be conditional on lifetimes or +/// user-implementable traits. +#[rustc_specialization_trait] +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe trait BoundedSize { + const UPPER_BOUND: Option = NonZeroUsize::new(1); +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl BoundedSize for Option {} +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl BoundedSize for option::IntoIter {} +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl BoundedSize for Result {} +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl BoundedSize for result::IntoIter {} +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl BoundedSize for Once {} +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl BoundedSize for OnceWith {} +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl BoundedSize for [T; N] { + const UPPER_BOUND: Option = NonZeroUsize::new(N); +} +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl BoundedSize for array::IntoIter { + const UPPER_BOUND: Option = NonZeroUsize::new(N); +} +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl BoundedSize for Filter { + const UPPER_BOUND: Option = I::UPPER_BOUND; +} +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl BoundedSize for FilterMap { + const UPPER_BOUND: Option = I::UPPER_BOUND; +} +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl BoundedSize for Map { + const UPPER_BOUND: Option = I::UPPER_BOUND; +} +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl BoundedSize for Copied { + const UPPER_BOUND: Option = I::UPPER_BOUND; +} +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl BoundedSize for Cloned { + const UPPER_BOUND: Option = I::UPPER_BOUND; +} + /// An iterator that flattens one level of nesting in an iterator of things /// that can be turned into iterators. /// @@ -289,6 +380,36 @@ where { } +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl InPlaceIterable for Flatten +where + I: InPlaceIterable + Iterator, + ::Item: IntoIterator + BoundedSize, +{ + const EXPAND_BY: Option = const { + match (I::EXPAND_BY, I::Item::UPPER_BOUND) { + (Some(m), Some(n)) => m.checked_mul(n), + _ => None, + } + }; + const MERGE_BY: Option = I::MERGE_BY; +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl SourceIter for Flatten +where + I: SourceIter + TrustedFused + Iterator, + ::Item: IntoIterator, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut I::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.inner.iter) } + } +} + #[stable(feature = "default_iters", since = "1.70.0")] impl Default for Flatten where diff --git a/library/core/src/iter/adapters/fuse.rs b/library/core/src/iter/adapters/fuse.rs index b1fa4f92117b1..3234cade50dc5 100644 --- a/library/core/src/iter/adapters/fuse.rs +++ b/library/core/src/iter/adapters/fuse.rs @@ -1,8 +1,9 @@ use crate::intrinsics; use crate::iter::adapters::zip::try_get_unchecked; +use crate::iter::adapters::SourceIter; use crate::iter::{ - DoubleEndedIterator, ExactSizeIterator, FusedIterator, TrustedLen, TrustedRandomAccess, - TrustedRandomAccessNoCoerce, + DoubleEndedIterator, ExactSizeIterator, FusedIterator, TrustedFused, TrustedLen, + TrustedRandomAccess, TrustedRandomAccessNoCoerce, }; use crate::ops::Try; @@ -29,6 +30,9 @@ impl Fuse { #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for Fuse where I: Iterator {} +#[unstable(issue = "none", feature = "trusted_fused")] +unsafe impl TrustedFused for Fuse where I: TrustedFused {} + // Any specialized implementation here is made internal // to avoid exposing default fns outside this trait. #[stable(feature = "rust1", since = "1.0.0")] @@ -418,6 +422,23 @@ where } } +// This is used by Flatten's SourceIter impl +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl SourceIter for Fuse +where + I: SourceIter + TrustedFused, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut I::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements. + // TrustedFused guarantees that we'll never encounter a case where `self.iter` would + // be set to None. + unsafe { SourceIter::as_inner(self.iter.as_mut().unwrap_unchecked()) } + } +} + #[inline] fn and_then_or_clear(opt: &mut Option, f: impl FnOnce(&mut T) -> Option) -> Option { let x = f(opt.as_mut()?); diff --git a/library/core/src/iter/adapters/inspect.rs b/library/core/src/iter/adapters/inspect.rs index 19839fdfe5bc3..fd2d830b693d1 100644 --- a/library/core/src/iter/adapters/inspect.rs +++ b/library/core/src/iter/adapters/inspect.rs @@ -1,5 +1,6 @@ use crate::fmt; -use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable}; +use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedFused}; +use crate::num::NonZeroUsize; use crate::ops::Try; /// An iterator that calls a function with a reference to each element before @@ -148,6 +149,9 @@ where #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for Inspect where F: FnMut(&I::Item) {} +#[unstable(issue = "none", feature = "trusted_fused")] +unsafe impl TrustedFused for Inspect {} + #[unstable(issue = "none", feature = "inplace_iteration")] unsafe impl SourceIter for Inspect where @@ -163,4 +167,7 @@ where } #[unstable(issue = "none", feature = "inplace_iteration")] -unsafe impl InPlaceIterable for Inspect where F: FnMut(&I::Item) {} +unsafe impl InPlaceIterable for Inspect { + const EXPAND_BY: Option = I::EXPAND_BY; + const MERGE_BY: Option = I::MERGE_BY; +} diff --git a/library/core/src/iter/adapters/map.rs b/library/core/src/iter/adapters/map.rs index 31d02a4da6ea5..e27fc7257f6cc 100644 --- a/library/core/src/iter/adapters/map.rs +++ b/library/core/src/iter/adapters/map.rs @@ -2,7 +2,8 @@ use crate::fmt; use crate::iter::adapters::{ zip::try_get_unchecked, SourceIter, TrustedRandomAccess, TrustedRandomAccessNoCoerce, }; -use crate::iter::{FusedIterator, InPlaceIterable, TrustedLen, UncheckedIterator}; +use crate::iter::{FusedIterator, InPlaceIterable, TrustedFused, TrustedLen, UncheckedIterator}; +use crate::num::NonZeroUsize; use crate::ops::Try; /// An iterator that maps the values of `iter` with `f`. @@ -179,6 +180,9 @@ where #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for Map where F: FnMut(I::Item) -> B {} +#[unstable(issue = "none", feature = "trusted_fused")] +unsafe impl TrustedFused for Map {} + #[unstable(feature = "trusted_len", issue = "37572")] unsafe impl TrustedLen for Map where @@ -228,4 +232,7 @@ where } #[unstable(issue = "none", feature = "inplace_iteration")] -unsafe impl InPlaceIterable for Map where F: FnMut(I::Item) -> B {} +unsafe impl InPlaceIterable for Map { + const EXPAND_BY: Option = I::EXPAND_BY; + const MERGE_BY: Option = I::MERGE_BY; +} diff --git a/library/core/src/iter/adapters/map_while.rs b/library/core/src/iter/adapters/map_while.rs index fbdeca4d4ee4f..bcae73cbe09cf 100644 --- a/library/core/src/iter/adapters/map_while.rs +++ b/library/core/src/iter/adapters/map_while.rs @@ -1,5 +1,6 @@ use crate::fmt; use crate::iter::{adapters::SourceIter, InPlaceIterable}; +use crate::num::NonZeroUsize; use crate::ops::{ControlFlow, Try}; /// An iterator that only accepts elements while `predicate` returns `Some(_)`. @@ -82,7 +83,7 @@ where } #[unstable(issue = "none", feature = "inplace_iteration")] -unsafe impl InPlaceIterable for MapWhile where - P: FnMut(I::Item) -> Option -{ +unsafe impl InPlaceIterable for MapWhile { + const EXPAND_BY: Option = I::EXPAND_BY; + const MERGE_BY: Option = I::MERGE_BY; } diff --git a/library/core/src/iter/adapters/mod.rs b/library/core/src/iter/adapters/mod.rs index 6f4fa7010f446..a691039c789c6 100644 --- a/library/core/src/iter/adapters/mod.rs +++ b/library/core/src/iter/adapters/mod.rs @@ -1,4 +1,5 @@ use crate::iter::{InPlaceIterable, Iterator}; +use crate::num::NonZeroUsize; use crate::ops::{ChangeOutputType, ControlFlow, FromResidual, Residual, Try}; mod array_chunks; @@ -119,8 +120,9 @@ pub unsafe trait SourceIter { /// /// # Safety /// - /// Implementations of must return the same mutable reference for their lifetime, unless + /// Implementations must return the same mutable reference for their lifetime, unless /// replaced by a caller. + /// /// Callers may only replace the reference when they stopped iteration and drop the /// iterator pipeline after extracting the source. /// @@ -228,7 +230,10 @@ where // in order to return `Some(_)`. Since `iter` has type `I: InPlaceIterable` it's // guaranteed that at least one item will be moved out from the underlying source. #[unstable(issue = "none", feature = "inplace_iteration")] -unsafe impl InPlaceIterable for GenericShunt<'_, I, R> where - I: Iterator> + InPlaceIterable +unsafe impl InPlaceIterable for GenericShunt<'_, I, R> +where + I: InPlaceIterable, { + const EXPAND_BY: Option = I::EXPAND_BY; + const MERGE_BY: Option = I::MERGE_BY; } diff --git a/library/core/src/iter/adapters/scan.rs b/library/core/src/iter/adapters/scan.rs index 62470512cc747..635bad199ff95 100644 --- a/library/core/src/iter/adapters/scan.rs +++ b/library/core/src/iter/adapters/scan.rs @@ -1,5 +1,6 @@ use crate::fmt; use crate::iter::{adapters::SourceIter, InPlaceIterable}; +use crate::num::NonZeroUsize; use crate::ops::{ControlFlow, Try}; /// An iterator to maintain state while iterating another iterator. @@ -92,7 +93,7 @@ where } #[unstable(issue = "none", feature = "inplace_iteration")] -unsafe impl InPlaceIterable for Scan where - F: FnMut(&mut St, I::Item) -> Option -{ +unsafe impl InPlaceIterable for Scan { + const EXPAND_BY: Option = I::EXPAND_BY; + const MERGE_BY: Option = I::MERGE_BY; } diff --git a/library/core/src/iter/adapters/skip.rs b/library/core/src/iter/adapters/skip.rs index 306338bc7cca0..e6c946e7f880d 100644 --- a/library/core/src/iter/adapters/skip.rs +++ b/library/core/src/iter/adapters/skip.rs @@ -1,4 +1,5 @@ use crate::intrinsics::unlikely; +use crate::iter::TrustedFused; use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable}; use crate::num::NonZeroUsize; use crate::ops::{ControlFlow, Try}; @@ -214,6 +215,9 @@ where #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for Skip where I: FusedIterator {} +#[unstable(issue = "none", feature = "trusted_fused")] +unsafe impl TrustedFused for Skip {} + #[unstable(issue = "none", feature = "inplace_iteration")] unsafe impl SourceIter for Skip where @@ -229,4 +233,7 @@ where } #[unstable(issue = "none", feature = "inplace_iteration")] -unsafe impl InPlaceIterable for Skip {} +unsafe impl InPlaceIterable for Skip { + const EXPAND_BY: Option = I::EXPAND_BY; + const MERGE_BY: Option = I::MERGE_BY; +} diff --git a/library/core/src/iter/adapters/skip_while.rs b/library/core/src/iter/adapters/skip_while.rs index f29661779c056..3a661973e5fe0 100644 --- a/library/core/src/iter/adapters/skip_while.rs +++ b/library/core/src/iter/adapters/skip_while.rs @@ -1,5 +1,6 @@ use crate::fmt; -use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable}; +use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedFused}; +use crate::num::NonZeroUsize; use crate::ops::Try; /// An iterator that rejects elements while `predicate` returns `true`. @@ -104,6 +105,9 @@ where { } +#[unstable(issue = "none", feature = "trusted_fused")] +unsafe impl TrustedFused for SkipWhile {} + #[unstable(issue = "none", feature = "inplace_iteration")] unsafe impl SourceIter for SkipWhile where @@ -119,7 +123,7 @@ where } #[unstable(issue = "none", feature = "inplace_iteration")] -unsafe impl InPlaceIterable for SkipWhile where - F: FnMut(&I::Item) -> bool -{ +unsafe impl InPlaceIterable for SkipWhile { + const EXPAND_BY: Option = I::EXPAND_BY; + const MERGE_BY: Option = I::MERGE_BY; } diff --git a/library/core/src/iter/adapters/take.rs b/library/core/src/iter/adapters/take.rs index c1d8cc4ff57bd..80e06066d28f9 100644 --- a/library/core/src/iter/adapters/take.rs +++ b/library/core/src/iter/adapters/take.rs @@ -1,6 +1,7 @@ use crate::cmp; use crate::iter::{ - adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedLen, TrustedRandomAccess, + adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedFused, TrustedLen, + TrustedRandomAccess, }; use crate::num::NonZeroUsize; use crate::ops::{ControlFlow, Try}; @@ -143,7 +144,10 @@ where } #[unstable(issue = "none", feature = "inplace_iteration")] -unsafe impl InPlaceIterable for Take {} +unsafe impl InPlaceIterable for Take { + const EXPAND_BY: Option = I::EXPAND_BY; + const MERGE_BY: Option = I::MERGE_BY; +} #[stable(feature = "double_ended_take_iterator", since = "1.38.0")] impl DoubleEndedIterator for Take @@ -241,6 +245,9 @@ impl ExactSizeIterator for Take where I: ExactSizeIterator {} #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for Take where I: FusedIterator {} +#[unstable(issue = "none", feature = "trusted_fused")] +unsafe impl TrustedFused for Take {} + #[unstable(feature = "trusted_len", issue = "37572")] unsafe impl TrustedLen for Take {} diff --git a/library/core/src/iter/adapters/take_while.rs b/library/core/src/iter/adapters/take_while.rs index ec66dc3aec360..e55d55a6d2377 100644 --- a/library/core/src/iter/adapters/take_while.rs +++ b/library/core/src/iter/adapters/take_while.rs @@ -1,5 +1,6 @@ use crate::fmt; -use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable}; +use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedFused}; +use crate::num::NonZeroUsize; use crate::ops::{ControlFlow, Try}; /// An iterator that only accepts elements while `predicate` returns `true`. @@ -105,6 +106,9 @@ where { } +#[unstable(issue = "none", feature = "trusted_fused")] +unsafe impl TrustedFused for TakeWhile {} + #[unstable(issue = "none", feature = "inplace_iteration")] unsafe impl SourceIter for TakeWhile where @@ -120,7 +124,7 @@ where } #[unstable(issue = "none", feature = "inplace_iteration")] -unsafe impl InPlaceIterable for TakeWhile where - F: FnMut(&I::Item) -> bool -{ +unsafe impl InPlaceIterable for TakeWhile { + const EXPAND_BY: Option = I::EXPAND_BY; + const MERGE_BY: Option = I::MERGE_BY; } diff --git a/library/core/src/iter/adapters/zip.rs b/library/core/src/iter/adapters/zip.rs index b6b0c90cb7d14..5993c1e304589 100644 --- a/library/core/src/iter/adapters/zip.rs +++ b/library/core/src/iter/adapters/zip.rs @@ -1,7 +1,8 @@ use crate::cmp; use crate::fmt::{self, Debug}; -use crate::iter::{DoubleEndedIterator, ExactSizeIterator, FusedIterator, Iterator}; +use crate::iter::{DoubleEndedIterator, ExactSizeIterator, FusedIterator, Iterator, TrustedFused}; use crate::iter::{InPlaceIterable, SourceIter, TrustedLen, UncheckedIterator}; +use crate::num::NonZeroUsize; /// An iterator that iterates two other iterators simultaneously. /// @@ -409,6 +410,14 @@ where { } +#[unstable(issue = "none", feature = "trusted_fused")] +unsafe impl TrustedFused for Zip +where + A: TrustedFused, + B: TrustedFused, +{ +} + #[unstable(feature = "trusted_len", issue = "37572")] unsafe impl TrustedLen for Zip where @@ -442,7 +451,10 @@ where // Since SourceIter forwards the left hand side we do the same here #[unstable(issue = "none", feature = "inplace_iteration")] -unsafe impl InPlaceIterable for Zip {} +unsafe impl InPlaceIterable for Zip { + const EXPAND_BY: Option = A::EXPAND_BY; + const MERGE_BY: Option = A::MERGE_BY; +} #[stable(feature = "rust1", since = "1.0.0")] impl Debug for Zip { diff --git a/library/core/src/iter/mod.rs b/library/core/src/iter/mod.rs index ca977d1ef8240..a727a97ed502d 100644 --- a/library/core/src/iter/mod.rs +++ b/library/core/src/iter/mod.rs @@ -417,6 +417,8 @@ pub use self::sources::{successors, Successors}; pub use self::traits::FusedIterator; #[unstable(issue = "none", feature = "inplace_iteration")] pub use self::traits::InPlaceIterable; +#[unstable(issue = "none", feature = "trusted_fused")] +pub use self::traits::TrustedFused; #[unstable(feature = "trusted_len", issue = "37572")] pub use self::traits::TrustedLen; #[unstable(feature = "trusted_step", issue = "85731")] diff --git a/library/core/src/iter/traits/marker.rs b/library/core/src/iter/traits/marker.rs index c21a2aac1c9a6..e7c1f195aacc6 100644 --- a/library/core/src/iter/traits/marker.rs +++ b/library/core/src/iter/traits/marker.rs @@ -1,4 +1,16 @@ use crate::iter::Step; +use crate::num::NonZeroUsize; + +/// Same as FusedIterator +/// +/// # Safety +/// +/// This is used for specialization. Therefore implementations must not +/// be lifetime-dependent. +#[unstable(issue = "none", feature = "trusted_fused")] +#[doc(hidden)] +#[rustc_specialization_trait] +pub unsafe trait TrustedFused {} /// An iterator that always continues to yield `None` when exhausted. /// @@ -14,6 +26,8 @@ use crate::iter::Step; /// [`Fuse`]: crate::iter::Fuse #[stable(feature = "fused", since = "1.26.0")] #[rustc_unsafe_specialization_marker] +// FIXME: this should be a #[marker] and have another blanket impl for T: TrustedFused +// but that ICEs iter::Fuse specializations. pub trait FusedIterator: Iterator {} #[stable(feature = "fused", since = "1.26.0")] @@ -71,7 +85,19 @@ unsafe impl TrustedLen for &mut I {} /// [`try_fold()`]: Iterator::try_fold #[unstable(issue = "none", feature = "inplace_iteration")] #[doc(hidden)] -pub unsafe trait InPlaceIterable: Iterator {} +#[rustc_specialization_trait] +pub unsafe trait InPlaceIterable { + /// The product of one-to-many item expansions that happen throughout the iterator pipeline. + /// E.g. [[u8; 4]; 4].iter().flatten().flatten() would have a `EXPAND_BY` of 16. + /// This is an upper bound, i.e. the transformations will produce at most this many items per + /// input. It's meant for layout calculations. + const EXPAND_BY: Option; + /// The product of many-to-one item reductions that happen throughout the iterator pipeline. + /// E.g. [u8].iter().array_chunks::<4>().array_chunks::<4>() would have a `MERGE_BY` of 16. + /// This is a lower bound, i.e. the transformations will consume at least this many items per + /// output. + const MERGE_BY: Option; +} /// A type that upholds all invariants of [`Step`]. /// diff --git a/library/core/src/iter/traits/mod.rs b/library/core/src/iter/traits/mod.rs index 41ea29e6a84d9..d4c9cc4b16037 100644 --- a/library/core/src/iter/traits/mod.rs +++ b/library/core/src/iter/traits/mod.rs @@ -18,6 +18,8 @@ pub use self::{ #[unstable(issue = "none", feature = "inplace_iteration")] pub use self::marker::InPlaceIterable; +#[unstable(issue = "none", feature = "trusted_fused")] +pub use self::marker::TrustedFused; #[unstable(feature = "trusted_step", issue = "85731")] pub use self::marker::TrustedStep;