diff options
| -rw-r--r-- | library/core/src/array/drain.rs | 33 | ||||
| -rw-r--r-- | library/core/src/array/mod.rs | 224 | ||||
| -rw-r--r-- | library/core/src/iter/adapters/array_chunks.rs | 14 | ||||
| -rw-r--r-- | library/core/src/iter/adapters/cloned.rs | 15 | ||||
| -rw-r--r-- | library/core/src/iter/adapters/map.rs | 15 | ||||
| -rw-r--r-- | library/core/src/iter/adapters/zip.rs | 9 | ||||
| -rw-r--r-- | library/core/src/iter/mod.rs | 1 | ||||
| -rw-r--r-- | library/core/src/iter/traits/mod.rs | 3 | ||||
| -rw-r--r-- | library/core/src/iter/traits/unchecked_iterator.rs | 36 | ||||
| -rw-r--r-- | library/core/src/ops/try_trait.rs | 9 | ||||
| -rw-r--r-- | library/core/src/slice/iter.rs | 4 | ||||
| -rw-r--r-- | library/core/src/slice/iter/macros.rs | 9 | ||||
| -rw-r--r-- | library/core/tests/iter/traits/iterator.rs | 3 | ||||
| -rw-r--r-- | tests/codegen/array-map.rs | 4 |
14 files changed, 237 insertions, 142 deletions
diff --git a/library/core/src/array/drain.rs b/library/core/src/array/drain.rs index 5ca93d54f87..5fadf907b62 100644 --- a/library/core/src/array/drain.rs +++ b/library/core/src/array/drain.rs @@ -1,11 +1,21 @@ -use crate::iter::TrustedLen; +use crate::iter::{TrustedLen, UncheckedIterator}; use crate::mem::ManuallyDrop; use crate::ptr::drop_in_place; use crate::slice; -// INVARIANT: It's ok to drop the remainder of the inner iterator. -pub(crate) struct Drain<'a, T>(slice::IterMut<'a, T>); - +/// A situationally-optimized version of `array.into_iter().for_each(func)`. +/// +/// [`crate::array::IntoIter`]s are great when you need an owned iterator, but +/// storing the entire array *inside* the iterator like that can sometimes +/// pessimize code. Notable, it can be more bytes than you really want to move +/// around, and because the array accesses index into it SRoA has a harder time +/// optimizing away the type than it does iterators that just hold a couple pointers. +/// +/// Thus this function exists, which gives a way to get *moved* access to the +/// elements of an array using a small iterator -- no bigger than a slice iterator. +/// +/// The function-taking-a-closure structure makes it safe, as it keeps callers +/// from looking at already-dropped elements. pub(crate) fn drain_array_with<T, R, const N: usize>( array: [T; N], func: impl for<'a> FnOnce(Drain<'a, T>) -> R, @@ -16,6 +26,11 @@ pub(crate) fn drain_array_with<T, R, const N: usize>( func(drain) } +/// See [`drain_array_with`] -- this is `pub(crate)` only so it's allowed to be +/// mentioned in the signature of that method. (Otherwise it hits `E0446`.) +// INVARIANT: It's ok to drop the remainder of the inner iterator. +pub(crate) struct Drain<'a, T>(slice::IterMut<'a, T>); + impl<T> Drop for Drain<'_, T> { fn drop(&mut self) { // SAFETY: By the type invariant, we're allowed to drop all these. @@ -49,3 +64,13 @@ impl<T> ExactSizeIterator for Drain<'_, T> { // SAFETY: This is a 1:1 wrapper for a slice iterator, which is also `TrustedLen`. unsafe impl<T> TrustedLen for Drain<'_, T> {} + +impl<T> UncheckedIterator for Drain<'_, T> { + unsafe fn next_unchecked(&mut self) -> T { + // SAFETY: `Drain` is 1:1 with the inner iterator, so if the caller promised + // that there's an element left, the inner iterator has one too. + let p: *const T = unsafe { self.0.next_unchecked() }; + // SAFETY: The iterator was already advanced, so we won't drop this later. + unsafe { p.read() } + } +} diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index 45ec68e6e7a..ae9f6e70f43 100644 --- a/library/core/src/array/mod.rs +++ b/library/core/src/array/mod.rs @@ -10,7 +10,7 @@ use crate::convert::{Infallible, TryFrom}; use crate::error::Error; use crate::fmt; use crate::hash::{self, Hash}; -use crate::iter::TrustedLen; +use crate::iter::UncheckedIterator; use crate::mem::{self, MaybeUninit}; use crate::ops::{ ChangeOutputType, ControlFlow, FromResidual, Index, IndexMut, NeverShortCircuit, Residual, Try, @@ -55,16 +55,11 @@ pub use iter::IntoIter; /// ``` #[inline] #[stable(feature = "array_from_fn", since = "1.63.0")] -pub fn from_fn<T, const N: usize, F>(mut cb: F) -> [T; N] +pub fn from_fn<T, const N: usize, F>(cb: F) -> [T; N] where F: FnMut(usize) -> T, { - let mut idx = 0; - [(); N].map(|_| { - let res = cb(idx); - idx += 1; - res - }) + try_from_fn(NeverShortCircuit::wrap_mut_1(cb)).0 } /// Creates an array `[T; N]` where each fallible array element `T` is returned by the `cb` call. @@ -104,9 +99,14 @@ where R: Try, R::Residual: Residual<[R::Output; N]>, { - // SAFETY: we know for certain that this iterator will yield exactly `N` - // items. - unsafe { try_collect_into_array_unchecked(&mut (0..N).map(cb)) } + let mut array = MaybeUninit::uninit_array::<N>(); + match try_from_fn_erased(&mut array, cb) { + ControlFlow::Break(r) => FromResidual::from_residual(r), + ControlFlow::Continue(()) => { + // SAFETY: All elements of the array were populated. + try { unsafe { MaybeUninit::array_assume_init(array) } } + } + } } /// Converts a reference to `T` into a reference to an array of length 1 (without copying). @@ -430,9 +430,7 @@ trait SpecArrayClone: Clone { impl<T: Clone> SpecArrayClone for T { #[inline] default fn clone<const N: usize>(array: &[T; N]) -> [T; N] { - // SAFETY: we know for certain that this iterator will yield exactly `N` - // items. - unsafe { collect_into_array_unchecked(&mut array.iter().cloned()) } + from_trusted_iterator(array.iter().cloned()) } } @@ -516,12 +514,7 @@ impl<T, const N: usize> [T; N] { where F: FnMut(T) -> U, { - drain_array_with(self, |iter| { - let mut iter = iter.map(f); - // SAFETY: we know for certain that this iterator will yield exactly `N` - // items. - unsafe { collect_into_array_unchecked(&mut iter) } - }) + self.try_map(NeverShortCircuit::wrap_mut_1(f)).0 } /// A fallible function `f` applied to each element on array `self` in order to @@ -558,12 +551,7 @@ impl<T, const N: usize> [T; N] { R: Try, R::Residual: Residual<[R::Output; N]>, { - drain_array_with(self, |iter| { - let mut iter = iter.map(f); - // SAFETY: we know for certain that this iterator will yield exactly `N` - // items. - unsafe { try_collect_into_array_unchecked(&mut iter) } - }) + drain_array_with(self, |iter| try_from_trusted_iterator(iter.map(f))) } /// 'Zips up' two arrays into a single array of pairs. @@ -585,12 +573,7 @@ impl<T, const N: usize> [T; N] { #[unstable(feature = "array_zip", issue = "80094")] pub fn zip<U>(self, rhs: [U; N]) -> [(T, U); N] { drain_array_with(self, |lhs| { - drain_array_with(rhs, |rhs| { - let mut iter = crate::iter::zip(lhs, rhs); - // SAFETY: we know for certain that this iterator will yield exactly `N` - // items. - unsafe { collect_into_array_unchecked(&mut iter) } - }) + drain_array_with(rhs, |rhs| from_trusted_iterator(crate::iter::zip(lhs, rhs))) }) } @@ -638,9 +621,7 @@ impl<T, const N: usize> [T; N] { /// ``` #[unstable(feature = "array_methods", issue = "76118")] pub fn each_ref(&self) -> [&T; N] { - // SAFETY: we know for certain that this iterator will yield exactly `N` - // items. - unsafe { collect_into_array_unchecked(&mut self.iter()) } + from_trusted_iterator(self.iter()) } /// Borrows each element mutably and returns an array of mutable references @@ -660,9 +641,7 @@ impl<T, const N: usize> [T; N] { /// ``` #[unstable(feature = "array_methods", issue = "76118")] pub fn each_mut(&mut self) -> [&mut T; N] { - // SAFETY: we know for certain that this iterator will yield exactly `N` - // items. - unsafe { collect_into_array_unchecked(&mut self.iter_mut()) } + from_trusted_iterator(self.iter_mut()) } /// Divides one array reference into two at an index. @@ -822,99 +801,71 @@ impl<T, const N: usize> [T; N] { } } -/// Pulls `N` items from `iter` and returns them as an array. If the iterator -/// yields fewer than `N` items, this function exhibits undefined behavior. +/// Populate an array from the first `N` elements of `iter` /// -/// # Safety +/// # Panics /// -/// It is up to the caller to guarantee that `iter` yields at least `N` items. -/// Violating this condition causes undefined behavior. -unsafe fn try_collect_into_array_unchecked<I, T, R, const N: usize>( - iter: &mut I, -) -> ChangeOutputType<I::Item, [T; N]> -where - // Note: `TrustedLen` here is somewhat of an experiment. This is just an - // internal function, so feel free to remove if this bound turns out to be a - // bad idea. In that case, remember to also remove the lower bound - // `debug_assert!` below! - I: Iterator + TrustedLen, - I::Item: Try<Output = T, Residual = R>, - R: Residual<[T; N]>, -{ - debug_assert!(N <= iter.size_hint().1.unwrap_or(usize::MAX)); - debug_assert!(N <= iter.size_hint().0); - - let mut array = MaybeUninit::uninit_array::<N>(); - let cf = try_collect_into_array_erased(iter, &mut array); - match cf { - ControlFlow::Break(r) => FromResidual::from_residual(r), - ControlFlow::Continue(initialized) => { - debug_assert_eq!(initialized, N); - // SAFETY: because of our function contract, all the elements - // must have been initialized. - let output = unsafe { MaybeUninit::array_assume_init(array) }; - Try::from_output(output) - } - } +/// If the iterator doesn't actually have enough items. +/// +/// By depending on `TrustedLen`, however, we can do that check up-front (where +/// it easily optimizes away) so it doesn't impact the loop that fills the array. +#[inline] +fn from_trusted_iterator<T, const N: usize>(iter: impl UncheckedIterator<Item = T>) -> [T; N] { + try_from_trusted_iterator(iter.map(NeverShortCircuit)).0 } -/// Infallible version of [`try_collect_into_array_unchecked`]. -unsafe fn collect_into_array_unchecked<I, const N: usize>(iter: &mut I) -> [I::Item; N] +#[inline] +fn try_from_trusted_iterator<T, R, const N: usize>( + iter: impl UncheckedIterator<Item = R>, +) -> ChangeOutputType<R, [T; N]> where - I: Iterator + TrustedLen, + R: Try<Output = T>, + R::Residual: Residual<[T; N]>, { - let mut map = iter.map(NeverShortCircuit); - - // SAFETY: The same safety considerations w.r.t. the iterator length - // apply for `try_collect_into_array_unchecked` as for - // `collect_into_array_unchecked` - match unsafe { try_collect_into_array_unchecked(&mut map) } { - NeverShortCircuit(array) => array, + assert!(iter.size_hint().0 >= N); + fn next<T>(mut iter: impl UncheckedIterator<Item = T>) -> impl FnMut(usize) -> T { + move |_| { + // SAFETY: We know that `from_fn` will call this at most N times, + // and we checked to ensure that we have at least that many items. + unsafe { iter.next_unchecked() } + } } + + try_from_fn(next(iter)) } -/// Rather than *returning* the array, this fills in a passed-in buffer. -/// If any of the iterator elements short-circuit, it drops everything in the -/// buffer and return the error. Otherwise it returns the number of items -/// which were initialized in the buffer. +/// Version of [`try_from_fn`] using a passed-in slice in order to avoid +/// needing to monomorphize for every array length. /// -/// (The caller is responsible for dropping those items on success, but not -/// doing that is just a leak, not UB, so this function is itself safe.) +/// This takes a generator rather than an iterator so that *at the type level* +/// it never needs to worry about running out of items. When combined with +/// an infallible `Try` type, that means the loop canonicalizes easily, allowing +/// it to optimize well. /// -/// This means less monomorphization, but more importantly it means that the -/// returned array doesn't need to be copied into the `Result`, since returning -/// the result seemed (2023-01) to cause in an extra `N + 1`-length `alloca` -/// even if it's always `unwrap_unchecked` later. +/// It would be *possible* to unify this and [`iter_next_chunk_erased`] into one +/// function that does the union of both things, but last time it was that way +/// it resulted in poor codegen from the "are there enough source items?" checks +/// not optimizing away. So if you give it a shot, make sure to watch what +/// happens in the codegen tests. #[inline] -fn try_collect_into_array_erased<I, T, R>( - iter: &mut I, +fn try_from_fn_erased<T, R>( buffer: &mut [MaybeUninit<T>], -) -> ControlFlow<R, usize> + mut generator: impl FnMut(usize) -> R, +) -> ControlFlow<R::Residual> where - I: Iterator, - I::Item: Try<Output = T, Residual = R>, + R: Try<Output = T>, { - let n = buffer.len(); let mut guard = Guard { array_mut: buffer, initialized: 0 }; - for _ in 0..n { - match iter.next() { - Some(item_rslt) => { - let item = item_rslt.branch()?; + while guard.initialized < guard.array_mut.len() { + let item = generator(guard.initialized).branch()?; - // SAFETY: `guard.initialized` starts at 0, which means push can be called - // at most `n` times, which this loop does. - unsafe { - guard.push_unchecked(item); - } - } - None => break, - } + // SAFETY: The loop condition ensures we have space to push the item + unsafe { guard.push_unchecked(item) }; } - let initialized = guard.initialized; mem::forget(guard); - ControlFlow::Continue(initialized) + ControlFlow::Continue(()) } /// Panic guard for incremental initialization of arrays. @@ -928,7 +879,7 @@ where /// /// To minimize indirection fields are still pub but callers should at least use /// `push_unchecked` to signal that something unsafe is going on. -pub(crate) struct Guard<'a, T> { +struct Guard<'a, T> { /// The array to be initialized. pub array_mut: &'a mut [MaybeUninit<T>], /// The number of items that have been initialized so far. @@ -960,7 +911,7 @@ impl<T> Drop for Guard<'_, T> { // SAFETY: this slice will contain only initialized objects. unsafe { crate::ptr::drop_in_place(MaybeUninit::slice_assume_init_mut( - &mut self.array_mut.get_unchecked_mut(..self.initialized), + self.array_mut.get_unchecked_mut(..self.initialized), )); } } @@ -982,17 +933,44 @@ impl<T> Drop for Guard<'_, T> { pub(crate) fn iter_next_chunk<T, const N: usize>( iter: &mut impl Iterator<Item = T>, ) -> Result<[T; N], IntoIter<T, N>> { - let mut map = iter.map(NeverShortCircuit); let mut array = MaybeUninit::uninit_array::<N>(); - let ControlFlow::Continue(initialized) = try_collect_into_array_erased(&mut map, &mut array); - if initialized == N { - // SAFETY: All elements of the array were populated. - let output = unsafe { MaybeUninit::array_assume_init(array) }; - Ok(output) - } else { - let alive = 0..initialized; - // SAFETY: `array` was initialized with exactly `initialized` - // number of elements. - return Err(unsafe { IntoIter::new_unchecked(array, alive) }); + let r = iter_next_chunk_erased(&mut array, iter); + match r { + Ok(()) => { + // SAFETY: All elements of `array` were populated. + Ok(unsafe { MaybeUninit::array_assume_init(array) }) + } + Err(initialized) => { + // SAFETY: Only the first `initialized` elements were populated + Err(unsafe { IntoIter::new_unchecked(array, 0..initialized) }) + } + } +} + +/// Version of [`iter_next_chunk`] using a passed-in slice in order to avoid +/// needing to monomorphize for every array length. +/// +/// Unfortunately this loop has two exit conditions, the buffer filling up +/// or the iterator running out of items, making it tend to optimize poorly. +#[inline] +fn iter_next_chunk_erased<T>( + buffer: &mut [MaybeUninit<T>], + iter: &mut impl Iterator<Item = T>, +) -> Result<(), usize> { + let mut guard = Guard { array_mut: buffer, initialized: 0 }; + while guard.initialized < guard.array_mut.len() { + let Some(item) = iter.next() else { + // Unlike `try_from_fn_erased`, we want to keep the partial results, + // so we need to defuse the guard instead of using `?`. + let initialized = guard.initialized; + mem::forget(guard); + return Err(initialized) + }; + + // SAFETY: The loop condition ensures we have space to push the item + unsafe { guard.push_unchecked(item) }; } + + mem::forget(guard); + Ok(()) } diff --git a/library/core/src/iter/adapters/array_chunks.rs b/library/core/src/iter/adapters/array_chunks.rs index af786609757..13719c727e9 100644 --- a/library/core/src/iter/adapters/array_chunks.rs +++ b/library/core/src/iter/adapters/array_chunks.rs @@ -1,6 +1,5 @@ use crate::array; use crate::iter::{ByRefSized, FusedIterator, Iterator, TrustedRandomAccessNoCoerce}; -use crate::mem::{self, MaybeUninit}; use crate::ops::{ControlFlow, NeverShortCircuit, Try}; /// An iterator over `N` elements of the iterator at a time. @@ -212,19 +211,14 @@ where let mut i = 0; // Use a while loop because (0..len).step_by(N) doesn't optimize well. while inner_len - i >= N { - let mut chunk = MaybeUninit::uninit_array(); - let mut guard = array::Guard { array_mut: &mut chunk, initialized: 0 }; - while guard.initialized < N { + let chunk = crate::array::from_fn(|local| { // SAFETY: The method consumes the iterator and the loop condition ensures that // all accesses are in bounds and only happen once. unsafe { - let idx = i + guard.initialized; - guard.push_unchecked(self.iter.__iterator_get_unchecked(idx)); + let idx = i + local; + self.iter.__iterator_get_unchecked(idx) } - } - mem::forget(guard); - // SAFETY: The loop above initialized all elements - let chunk = unsafe { MaybeUninit::array_assume_init(chunk) }; + }); accum = f(accum, chunk); i += N; } diff --git a/library/core/src/iter/adapters/cloned.rs b/library/core/src/iter/adapters/cloned.rs index aba24a79dcf..914ff86c1a9 100644 --- a/library/core/src/iter/adapters/cloned.rs +++ b/library/core/src/iter/adapters/cloned.rs @@ -1,7 +1,7 @@ use crate::iter::adapters::{ zip::try_get_unchecked, TrustedRandomAccess, TrustedRandomAccessNoCoerce, }; -use crate::iter::{FusedIterator, TrustedLen}; +use crate::iter::{FusedIterator, TrustedLen, UncheckedIterator}; use crate::ops::Try; /// An iterator that clones the elements of an underlying iterator. @@ -140,3 +140,16 @@ where T: Clone, { } + +impl<'a, I, T: 'a> UncheckedIterator for Cloned<I> +where + I: UncheckedIterator<Item = &'a T>, + T: Clone, +{ + unsafe fn next_unchecked(&mut self) -> T { + // SAFETY: `Cloned` is 1:1 with the inner iterator, so if the caller promised + // that there's an element left, the inner iterator has one too. + let item = unsafe { self.it.next_unchecked() }; + item.clone() + } +} diff --git a/library/core/src/iter/adapters/map.rs b/library/core/src/iter/adapters/map.rs index 9e25dbe462c..31d02a4da6e 100644 --- a/library/core/src/iter/adapters/map.rs +++ b/library/core/src/iter/adapters/map.rs @@ -2,7 +2,7 @@ use crate::fmt; use crate::iter::adapters::{ zip::try_get_unchecked, SourceIter, TrustedRandomAccess, TrustedRandomAccessNoCoerce, }; -use crate::iter::{FusedIterator, InPlaceIterable, TrustedLen}; +use crate::iter::{FusedIterator, InPlaceIterable, TrustedLen, UncheckedIterator}; use crate::ops::Try; /// An iterator that maps the values of `iter` with `f`. @@ -187,6 +187,19 @@ where { } +impl<B, I, F> UncheckedIterator for Map<I, F> +where + I: UncheckedIterator, + F: FnMut(I::Item) -> B, +{ + unsafe fn next_unchecked(&mut self) -> B { + // SAFETY: `Map` is 1:1 with the inner iterator, so if the caller promised + // that there's an element left, the inner iterator has one too. + let item = unsafe { self.iter.next_unchecked() }; + (self.f)(item) + } +} + #[doc(hidden)] #[unstable(feature = "trusted_random_access", issue = "none")] unsafe impl<I, F> TrustedRandomAccess for Map<I, F> where I: TrustedRandomAccess {} diff --git a/library/core/src/iter/adapters/zip.rs b/library/core/src/iter/adapters/zip.rs index 8153c8cfef1..b6b0c90cb7d 100644 --- a/library/core/src/iter/adapters/zip.rs +++ b/library/core/src/iter/adapters/zip.rs @@ -1,7 +1,7 @@ use crate::cmp; use crate::fmt::{self, Debug}; use crate::iter::{DoubleEndedIterator, ExactSizeIterator, FusedIterator, Iterator}; -use crate::iter::{InPlaceIterable, SourceIter, TrustedLen}; +use crate::iter::{InPlaceIterable, SourceIter, TrustedLen, UncheckedIterator}; /// An iterator that iterates two other iterators simultaneously. /// @@ -417,6 +417,13 @@ where { } +impl<A, B> UncheckedIterator for Zip<A, B> +where + A: UncheckedIterator, + B: UncheckedIterator, +{ +} + // Arbitrarily selects the left side of the zip iteration as extractable "source" // it would require negative trait bounds to be able to try both #[unstable(issue = "none", feature = "inplace_iteration")] diff --git a/library/core/src/iter/mod.rs b/library/core/src/iter/mod.rs index 00f57fbcc61..156b925de77 100644 --- a/library/core/src/iter/mod.rs +++ b/library/core/src/iter/mod.rs @@ -450,6 +450,7 @@ pub use self::adapters::{ pub use self::adapters::{Intersperse, IntersperseWith}; pub(crate) use self::adapters::try_process; +pub(crate) use self::traits::UncheckedIterator; mod adapters; mod range; diff --git a/library/core/src/iter/traits/mod.rs b/library/core/src/iter/traits/mod.rs index ed0fb634dbf..41ea29e6a84 100644 --- a/library/core/src/iter/traits/mod.rs +++ b/library/core/src/iter/traits/mod.rs @@ -4,6 +4,7 @@ mod double_ended; mod exact_size; mod iterator; mod marker; +mod unchecked_iterator; #[stable(feature = "rust1", since = "1.0.0")] pub use self::{ @@ -19,3 +20,5 @@ pub use self::{ pub use self::marker::InPlaceIterable; #[unstable(feature = "trusted_step", issue = "85731")] pub use self::marker::TrustedStep; + +pub(crate) use self::unchecked_iterator::UncheckedIterator; diff --git a/library/core/src/iter/traits/unchecked_iterator.rs b/library/core/src/iter/traits/unchecked_iterator.rs new file mode 100644 index 00000000000..ae4bfcad4e6 --- /dev/null +++ b/library/core/src/iter/traits/unchecked_iterator.rs @@ -0,0 +1,36 @@ +use crate::iter::TrustedLen; + +/// [`TrustedLen`] cannot have methods, so this allows augmenting it. +/// +/// It currently requires `TrustedLen` because it's unclear whether it's +/// reasonably possible to depend on the `size_hint` of anything else. +pub(crate) trait UncheckedIterator: TrustedLen { + /// Gets the next item from a non-empty iterator. + /// + /// Because there's always a value to return, that means it can return + /// the `Item` type directly, without wrapping it in an `Option`. + /// + /// # Safety + /// + /// This can only be called if `size_hint().0 != 0`, guaranteeing that + /// there's at least one item available. + /// + /// Otherwise (aka when `size_hint().1 == Some(0)`), this is UB. + /// + /// # Note to Implementers + /// + /// This has a default implementation using [`Option::unwrap_unchecked`]. + /// That's probably sufficient if your `next` *always* returns `Some`, + /// such as for infinite iterators. In more complicated situations, however, + /// sometimes there can still be `insertvalue`/`assume`/`extractvalue` + /// instructions remaining in the IR from the `Option` handling, at which + /// point you might want to implement this manually instead. + #[unstable(feature = "trusted_len_next_unchecked", issue = "37572")] + #[inline] + unsafe fn next_unchecked(&mut self) -> Self::Item { + let opt = self.next(); + // SAFETY: The caller promised that we're not empty, and + // `Self: TrustedLen` so we can actually trust the `size_hint`. + unsafe { opt.unwrap_unchecked() } + } +} diff --git a/library/core/src/ops/try_trait.rs b/library/core/src/ops/try_trait.rs index 9108fc63045..86aa1e4fd20 100644 --- a/library/core/src/ops/try_trait.rs +++ b/library/core/src/ops/try_trait.rs @@ -379,6 +379,15 @@ pub(crate) type ChangeOutputType<T, V> = <<T as Try>::Residual as Residual<V>>:: pub(crate) struct NeverShortCircuit<T>(pub T); impl<T> NeverShortCircuit<T> { + /// Wraps a unary function to produce one that wraps the output into a `NeverShortCircuit`. + /// + /// This is useful for implementing infallible functions in terms of the `try_` ones, + /// without accidentally capturing extra generic parameters in a closure. + #[inline] + pub fn wrap_mut_1<A>(mut f: impl FnMut(A) -> T) -> impl FnMut(A) -> NeverShortCircuit<T> { + move |a| NeverShortCircuit(f(a)) + } + #[inline] pub fn wrap_mut_2<A, B>( mut f: impl ~const FnMut(A, B) -> T, diff --git a/library/core/src/slice/iter.rs b/library/core/src/slice/iter.rs index 90ab43d1289..c4317799bcc 100644 --- a/library/core/src/slice/iter.rs +++ b/library/core/src/slice/iter.rs @@ -7,7 +7,9 @@ use crate::cmp; use crate::cmp::Ordering; use crate::fmt; use crate::intrinsics::assume; -use crate::iter::{FusedIterator, TrustedLen, TrustedRandomAccess, TrustedRandomAccessNoCoerce}; +use crate::iter::{ + FusedIterator, TrustedLen, TrustedRandomAccess, TrustedRandomAccessNoCoerce, UncheckedIterator, +}; use crate::marker::{PhantomData, Send, Sized, Sync}; use crate::mem::{self, SizedTypeProperties}; use crate::num::NonZeroUsize; diff --git a/library/core/src/slice/iter/macros.rs b/library/core/src/slice/iter/macros.rs index 0fd57b197aa..89b92a7d597 100644 --- a/library/core/src/slice/iter/macros.rs +++ b/library/core/src/slice/iter/macros.rs @@ -384,6 +384,15 @@ macro_rules! iterator { #[unstable(feature = "trusted_len", issue = "37572")] unsafe impl<T> TrustedLen for $name<'_, T> {} + + impl<'a, T> UncheckedIterator for $name<'a, T> { + unsafe fn next_unchecked(&mut self) -> $elem { + // SAFETY: The caller promised there's at least one more item. + unsafe { + next_unchecked!(self) + } + } + } } } diff --git a/library/core/tests/iter/traits/iterator.rs b/library/core/tests/iter/traits/iterator.rs index 37345c1d381..62566a9502d 100644 --- a/library/core/tests/iter/traits/iterator.rs +++ b/library/core/tests/iter/traits/iterator.rs @@ -582,6 +582,9 @@ fn test_next_chunk() { assert_eq!(it.next_chunk().unwrap(), []); assert_eq!(it.next_chunk().unwrap(), [4, 5, 6, 7, 8, 9]); assert_eq!(it.next_chunk::<4>().unwrap_err().as_slice(), &[10, 11]); + + let mut it = std::iter::repeat_with(|| panic!()); + assert_eq!(it.next_chunk::<0>().unwrap(), []); } // just tests by whether or not this compiles diff --git a/tests/codegen/array-map.rs b/tests/codegen/array-map.rs index 1154659eea5..9298e89e397 100644 --- a/tests/codegen/array-map.rs +++ b/tests/codegen/array-map.rs @@ -43,5 +43,7 @@ pub fn long_integer_map(x: [u32; 64]) -> [u32; 64] { // CHECK-NEXT: alloca [64 x i32] // CHECK-NEXT: alloca %"core::mem::manually_drop::ManuallyDrop<[u32; 64]>" // CHECK-NOT: alloca - x.map(|x| 2 * x + 1) + // CHECK: mul <{{[0-9]+}} x i32> + // CHECK: add <{{[0-9]+}} x i32> + x.map(|x| 13 * x + 7) } |