diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index 6cca2e6358b63..8a017bb88ee6d 100644 --- a/library/core/src/array/mod.rs +++ b/library/core/src/array/mod.rs @@ -964,6 +964,53 @@ impl const Drop for Guard<'_, T> { } } +/// Panic guard for incremental initialization of arrays from the back. +/// +/// Elements of the array are populated starting from the end towards the beginning. +/// Disarm the guard with `mem::forget` once the array has been fully initialized. +/// +/// # Safety +/// +/// All write accesses to this structure are unsafe and must maintain a correct +/// count of `initialized` elements. +struct GuardBack<'a, T> { + /// The array to be initialized (will be filled from the end). + pub array_mut: &'a mut [MaybeUninit], + /// The number of items that have been initialized so far. + pub initialized: usize, +} + +impl GuardBack<'_, T> { + /// Adds an item to the array and updates the initialized item counter. + /// + /// # Safety + /// + /// No more than N elements must be initialized. + #[inline] + pub(crate) unsafe fn push_unchecked(&mut self, item: T) { + // SAFETY: If `initialized` was correct before and the caller does not + // invoke this method more than N times, then writes will be in-bounds + // and slots will not be initialized more than once. + unsafe { + self.initialized = self.initialized.unchecked_add(1); + let index = self.array_mut.len().unchecked_sub(self.initialized); + self.array_mut.get_unchecked_mut(index).write(item); + } + } +} + +impl Drop for GuardBack<'_, T> { + #[inline] + fn drop(&mut self) { + debug_assert!(self.initialized <= self.array_mut.len()); + let len = self.array_mut.len(); + // SAFETY: this slice will contain only initialized objects. + unsafe { + self.array_mut.get_unchecked_mut(len - self.initialized..len).assume_init_drop(); + } + } +} + /// Pulls `N` items from `iter` and returns them as an array. If the iterator /// yields fewer than `N` items, `Err` is returned containing an iterator over /// the already yielded items. @@ -1022,3 +1069,56 @@ fn iter_next_chunk_erased( mem::forget(guard); Ok(()) } + +/// Pulls `N` items from the back of `iter` and returns them as an array. +/// If the iterator yields fewer than `N` items, `Err` is returned containing +/// an iterator over the already yielded items. +/// +/// Since the iterator is passed as a mutable reference and this function calls +/// `next_back` at most `N` times, the iterator can still be used afterwards to +/// retrieve the remaining items. +/// +/// If `iter.next_back()` panics, all items already yielded by the iterator are +/// dropped. +/// +/// Used for [`DoubleEndedIterator::next_chunk_back`]. +#[inline] +pub(crate) fn iter_next_chunk_back( + iter: &mut impl DoubleEndedIterator, +) -> Result<[T; N], IntoIter> { + let mut array = [const { MaybeUninit::uninit() }; N]; + let r = iter_next_chunk_back_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 last `initialized` elements were populated + Err(unsafe { IntoIter::new_unchecked(array, N - initialized..N) }) + } + } +} + +/// Version of [`iter_next_chunk_back`] using a passed-in slice. +#[inline] +fn iter_next_chunk_back_erased( + buffer: &mut [MaybeUninit], + iter: &mut impl DoubleEndedIterator, +) -> Result<(), usize> { + // if `Iterator::next_back` panics, this guard will drop already initialized items + let mut guard = GuardBack { array_mut: buffer, initialized: 0 }; + while guard.initialized < guard.array_mut.len() { + let Some(item) = iter.next_back() else { + 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 7c003cff10c7b..16cc515bb8822 100644 --- a/library/core/src/iter/adapters/array_chunks.rs +++ b/library/core/src/iter/adapters/array_chunks.rs @@ -1,8 +1,6 @@ use crate::array; use crate::iter::adapters::SourceIter; -use crate::iter::{ - ByRefSized, FusedIterator, InPlaceIterable, TrustedFused, TrustedRandomAccessNoCoerce, -}; +use crate::iter::{FusedIterator, InPlaceIterable, TrustedFused, TrustedRandomAccessNoCoerce}; use crate::num::NonZero; use crate::ops::{ControlFlow, NeverShortCircuit, Try}; @@ -128,15 +126,11 @@ where self.next_back_remainder(); let mut acc = init; - let mut iter = ByRefSized(&mut self.iter).rev(); // NB remainder is handled by `next_back_remainder`, so - // `next_chunk` can't return `Err` with non-empty remainder + // `next_chunk_back` can't return `Err` with non-empty remainder // (assuming correct `I as ExactSizeIterator` impl). - while let Ok(mut chunk) = iter.next_chunk() { - // FIXME: do not do double reverse - // (we could instead add `next_chunk_back` for example) - chunk.reverse(); + while let Ok(chunk) = self.iter.next_chunk_back() { acc = f(acc, chunk)? } diff --git a/library/core/src/iter/traits/double_ended.rs b/library/core/src/iter/traits/double_ended.rs index 9f7ac7da2dbda..1707d11f7939b 100644 --- a/library/core/src/iter/traits/double_ended.rs +++ b/library/core/src/iter/traits/double_ended.rs @@ -1,3 +1,4 @@ +use crate::array; use crate::num::NonZero; use crate::ops::{ControlFlow, Try}; @@ -93,6 +94,20 @@ pub trait DoubleEndedIterator: Iterator { #[stable(feature = "rust1", since = "1.0.0")] fn next_back(&mut self) -> Option; + /// Pulls `N` items from the back of the iterator and returns them as an array. + /// + /// See [`Iterator::next_chunk`] for more details. + #[inline] + #[unstable(feature = "iter_next_chunk", issue = "98326")] + fn next_chunk_back( + &mut self, + ) -> Result<[Self::Item; N], array::IntoIter> + where + Self: Sized, + { + crate::array::iter_next_chunk_back(self) + } + /// Advances the iterator from the back by `n` elements. /// /// `advance_back_by` is the reverse version of [`advance_by`]. This method will