Optimize ArrayChunks by adding next_chunk_back to DoubleEndedIterator

Author: vinDelphini

library/core/src/array/mod.rs

@@ -964,16 +964,46 @@ impl<T: [const] Destruct> const Drop for Guard<'_, T> {
     }
 }
 
-/// 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.
-///
-/// Since the iterator is passed as a mutable reference and this function calls
-/// `next` at most `N` times, the iterator can still be used afterwards to
-/// retrieve the remaining items.
-///
-/// If `iter.next()` panics, all items already yielded by the iterator are
-/// dropped.
+/// Panic guard for incremental initialization of arrays from the back.
+struct GuardBack<'a, T> {
+    /// The array to be initialized.
+    pub array_mut: &'a mut [MaybeUninit<T>],
+    /// The number of items that have been initialized so far.
+    pub initialized: usize,
+}
+
+impl<T> 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<T: Destruct> 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`.
 ///
 /// Used for [`Iterator::next_chunk`].
 #[inline]
@@ -994,11 +1024,7 @@ pub(crate) fn iter_next_chunk<T, const N: usize>(
     }
 }
 
-/// 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.
+/// Version of [`iter_next_chunk`] using a passed-in slice.
 #[inline]
 fn iter_next_chunk_erased<T>(
     buffer: &mut [MaybeUninit<T>],
@@ -1022,3 +1048,47 @@ fn iter_next_chunk_erased<T>(
     mem::forget(guard);
     Ok(())
 }
+
+/// Pulls `N` items from the back of `iter`.
+///
+/// Used for [`DoubleEndedIterator::next_chunk_back`].
+#[inline]
+pub(crate) fn iter_next_chunk_back<T, const N: usize>(
+    iter: &mut impl DoubleEndedIterator<Item = T>,
+) -> Result<[T; N], IntoIter<T, N>> {
+    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<T>(
+    buffer: &mut [MaybeUninit<T>],
+    iter: &mut impl DoubleEndedIterator<Item = T>,
+) -> 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(())
+}

library/core/src/iter/adapters/array_chunks.rs

@@ -1,7 +1,7 @@
 use crate::array;
 use crate::iter::adapters::SourceIter;
 use crate::iter::{
-    ByRefSized, FusedIterator, InPlaceIterable, TrustedFused, TrustedRandomAccessNoCoerce,
+    FusedIterator, InPlaceIterable, TrustedFused, TrustedRandomAccessNoCoerce,
 };
 use crate::num::NonZero;
 use crate::ops::{ControlFlow, NeverShortCircuit, Try};
@@ -128,15 +128,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)?
         }
 

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<Self::Item>;
 
+    /// 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<const N: usize>(
+        &mut self,
+    ) -> Result<[Self::Item; N], array::IntoIter<Self::Item, N>>
+    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