Implement DoubleEndedIterator::next_chunk_back

library/alloc/src/lib.rs

@@ -136,6 +136,7 @@
 #![feature(legacy_receiver_trait)]
 #![feature(likely_unlikely)]
 #![feature(local_waker)]
+#![feature(maybe_uninit_array_assume_init)]
 #![feature(maybe_uninit_uninit_array_transpose)]
 #![feature(panic_internals)]
 #![feature(pattern)]

library/alloc/src/vec/into_iter.rs

@@ -445,6 +445,47 @@ impl<T, A: Allocator> DoubleEndedIterator for IntoIter<T, A> {
         }
     }
 
+    #[inline]
+    fn next_chunk_back<const N: usize>(&mut self) -> Result<[T; N], core::array::IntoIter<T, N>> {
+        let mut raw_ary = [const { MaybeUninit::uninit() }; N];
+
+        let len = self.len();
+
+        if T::IS_ZST {
+            if len < N {
+                self.forget_remaining_elements();
+                // Safety: ZSTs can be conjured ex nihilo, only the amount has to be correct
+                return Err(unsafe { array::IntoIter::new_unchecked(raw_ary, N - len..N) });
+            }
+
+            self.end = self.end.wrapping_byte_sub(N);
+            // Safety: ditto
+            return Ok(unsafe { MaybeUninit::array_assume_init(raw_ary) });
+        }
+
+        if len < N {
+            // Safety: `len` indicates that this many elements are available and we just checked that
+            // it fits into the array.
+            unsafe {
+                ptr::copy_nonoverlapping(self.ptr.as_ptr(), raw_ary.as_mut_ptr() as *mut T, len);
+                self.forget_remaining_elements();
+                return Err(array::IntoIter::new_unchecked(raw_ary, 0..len));
+            }
+        }
+
+        // Safety: `len` is larger than the array size. Copy a fixed amount here to fully initialize
+        // the array.
+        unsafe {
+            ptr::copy_nonoverlapping(
+                self.ptr.add(len - N).as_ptr(),
+                raw_ary.as_mut_ptr() as *mut T,
+                N,
+            );
+            self.end = self.end.sub(N);
+            Ok(MaybeUninit::array_assume_init(raw_ary))
+        }
+    }
+
     #[inline]
     fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
         let step_size = self.len().min(n);

library/alloctests/tests/vec.rs

@@ -1027,6 +1027,15 @@ fn test_into_iter_next_chunk() {
     assert_eq!(iter.next_chunk::<4>().unwrap_err().as_slice(), &[]); // N is explicitly 4
 }
 
+#[test]
+fn test_into_iter_next_chunk_back() {
+    let mut iter = b"lorem".to_vec().into_iter();
+
+    assert_eq!(iter.next_chunk_back().unwrap(), [b'e', b'm']); // N is inferred as 2
+    assert_eq!(iter.next_chunk_back().unwrap(), [b'l', b'o', b'r']); // N is inferred as 3
+    assert_eq!(iter.next_chunk_back::<4>().unwrap_err().as_slice(), &[]); // N is explicitly 4
+}
+
 #[test]
 fn test_into_iter_clone() {
     fn iter_equal<I: Iterator<Item = i32>>(it: I, slice: &[i32]) {
@@ -1131,6 +1140,14 @@ fn test_into_iter_zst() {
     let mut it = vec![C, C].into_iter();
     it.next_chunk::<4>().unwrap_err();
     drop(it);
+
+    let mut it = vec![C, C].into_iter();
+    it.next_chunk_back::<1>().unwrap();
+    drop(it);
+
+    let mut it = vec![C, C].into_iter();
+    it.next_chunk_back::<4>().unwrap_err();
+    drop(it);
 }
 
 #[test]

library/core/src/array/mod.rs

@@ -969,6 +969,54 @@ const impl<T: [const] Destruct> 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<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 {
+            let offset = self.initialized.unchecked_add(1);
+            let index = self.array_mut.len().unchecked_sub(offset);
+            self.array_mut.get_unchecked_mut(index).write(item);
+            self.initialized = offset;
+        }
+    }
+}
+
+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` 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.
@@ -1076,3 +1124,109 @@ const fn iter_next_chunk_erased<T>(
     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<T, const N: usize>(
+    iter: &mut impl DoubleEndedIterator<Item = T>,
+) -> Result<[T; N], IntoIter<T, N>> {
+    iter.spec_next_chunk_back()
+}
+
+pub(crate) trait SpecNextChunkBack<T, const N: usize>:
+    DoubleEndedIterator<Item = T>
+{
+    fn spec_next_chunk_back(&mut self) -> Result<[T; N], IntoIter<T, N>>;
+}
+
+impl<I: DoubleEndedIterator<Item = T>, T, const N: usize> SpecNextChunkBack<T, N> for I {
+    #[inline]
+    default fn spec_next_chunk_back(&mut self) -> Result<[T; N], IntoIter<T, N>> {
+        let mut array = [const { MaybeUninit::uninit() }; N];
+        let r = iter_next_chunk_back_erased(&mut array, self);
+        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) })
+            }
+        }
+    }
+}
+
+impl<I: DoubleEndedIterator<Item = T> + TrustedLen, T, const N: usize> SpecNextChunkBack<T, N>
+    for I
+{
+    fn spec_next_chunk_back(&mut self) -> Result<[T; N], IntoIter<T, N>> {
+        let len = (*self).size_hint().0;
+        let mut array = [const { MaybeUninit::uninit() }; N];
+        if len < N {
+            // SAFETY: `TrustedLen`, an unsafe trait, requires that i can get len items out of it.
+            unsafe { write_back(&mut array, self, len) };
+            // SAFETY: Only the last `len` elements were populated
+            Err(unsafe { IntoIter::new_unchecked(array, N - len..N) })
+        } else {
+            // SAFETY: `TrustedLen`, an unsafe trait, requires that i can get N items out of it.
+            unsafe { write_back(&mut array, self, N) };
+            // SAFETY: All N items were populated
+            Ok(unsafe { MaybeUninit::array_assume_init(array) })
+        }
+    }
+}
+
+// SAFETY: `from` must have len items, and len items must be < N.
+unsafe fn write_back<T, const N: usize>(
+    to: &mut [MaybeUninit<T>; N],
+    from: &mut impl DoubleEndedIterator<Item = T>,
+    len: usize,
+) {
+    let mut guard = GuardBack { array_mut: to, initialized: 0 };
+    while guard.initialized < len {
+        // SAFETY: caller has guaranteed, from has len items.
+        let item = unsafe { from.next_back().unwrap_unchecked() };
+        // SAFETY: guard.initialized < len < N
+        unsafe { guard.push_unchecked(item) };
+    }
+    crate::mem::forget(guard);
+}
+
+/// Version of [`iter_next_chunk_back`] 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_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/traits/double_ended.rs

@@ -1,3 +1,4 @@
+use crate::array;
 use crate::marker::Destruct;
 use crate::num::NonZero;
 use crate::ops::{ControlFlow, Try};
@@ -95,6 +96,53 @@ pub const trait DoubleEndedIterator: [const] Iterator {
     #[stable(feature = "rust1", since = "1.0.0")]
     fn next_back(&mut self) -> Option<Self::Item>;
 
+    /// Advances from the back of the iterator and returns an array containing the next
+    /// `N` values in sequence.
+    ///
+    /// If there are not enough elements to fill the array then `Err` is returned
+    /// containing an iterator over the remaining elements.
+    ///
+    /// Note: This is not equivalent to doing `iter.rev().next_chunk()` as this method
+    /// takes elements from the back of the iterator and preserves the order that the
+    /// elements were seen in the original iterator.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// #![feature(iter_next_chunk)]
+    ///
+    /// let mut iter = "lorem".chars();
+    ///
+    /// assert_eq!(iter.next_chunk_back().unwrap(), ['e', 'm']);              // N is inferred as 2
+    /// assert_eq!(iter.next_chunk_back().unwrap(), ['l', 'o', 'r']);         // N is inferred as 3
+    /// assert_eq!(iter.next_chunk_back::<4>().unwrap_err().as_slice(), &[]); // N is explicitly 4
+    /// ```
+    ///
+    /// Split a string and get the last three items in sequence.
+    ///
+    /// ```
+    /// #![feature(iter_next_chunk)]
+    ///
+    /// let quote = "not all those who wander are lost";
+    /// let [first, second, third] = quote.split_whitespace().next_chunk_back().unwrap();
+    /// assert_eq!(first, "wander");
+    /// assert_eq!(second, "are");
+    /// assert_eq!(third, "lost");
+    /// ```
+    #[inline]
+    #[unstable(feature = "iter_next_chunk", issue = "98326")]
+    #[rustc_non_const_trait_method]
+    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

library/coretests/tests/iter/adapters/filter.rs

@@ -63,3 +63,15 @@ fn test_next_chunk_does_not_leak() {
         assert_eq!(Rc::strong_count(w), 1);
     }
 }
+
+#[test]
+fn test_next_chunk_back_does_not_leak() {
+    let drop_witness: [_; 5] = std::array::from_fn(|_| Rc::new(()));
+
+    let v = (0..5).map(|i| drop_witness[i].clone()).collect::<Vec<_>>();
+    let _ = v.into_iter().filter(|_| false).next_chunk_back::<1>();
+
+    for ref w in drop_witness {
+        assert_eq!(Rc::strong_count(w), 1);
+    }
+}

library/coretests/tests/iter/traits/iterator.rs

@@ -619,6 +619,18 @@ fn test_next_chunk() {
     assert_eq!(it.next_chunk::<0>().unwrap(), []);
 }
 
+#[test]
+fn test_next_chunk_back() {
+    let mut it = 0..12;
+    assert_eq!(it.next_chunk_back().unwrap(), [8, 9, 10, 11]);
+    assert_eq!(it.next_chunk_back().unwrap(), []);
+    assert_eq!(it.next_chunk_back().unwrap(), [2, 3, 4, 5, 6, 7]);
+    assert_eq!(it.next_chunk_back::<4>().unwrap_err().as_slice(), &[0, 1]);
+
+    let mut it = std::iter::once_with(|| panic!());
+    assert_eq!(it.next_chunk_back::<0>().unwrap(), []);
+}
+
 #[test]
 fn test_next_chunk_untrusted() {
     struct Untrusted<I: Iterator>(I);
@@ -636,6 +648,28 @@ fn test_next_chunk_untrusted() {
     assert_eq!(it.next_chunk::<4>().unwrap_err().as_slice(), &[10, 11]);
 }
 
+#[test]
+fn test_next_chunk_back_untrusted() {
+    struct Untrusted<I: Iterator>(I);
+    impl<I: Iterator> Iterator for Untrusted<I> {
+        type Item = I::Item;
+
+        fn next(&mut self) -> Option<Self::Item> {
+            self.0.next()
+        }
+    }
+    impl<I: DoubleEndedIterator> DoubleEndedIterator for Untrusted<I> {
+        fn next_back(&mut self) -> Option<Self::Item> {
+            self.0.next_back()
+        }
+    }
+    let mut it = Untrusted(0..12);
+    assert_eq!(it.next_chunk_back().unwrap(), [8, 9, 10, 11]);
+    assert_eq!(it.next_chunk_back().unwrap(), []);
+    assert_eq!(it.next_chunk_back().unwrap(), [2, 3, 4, 5, 6, 7]);
+    assert_eq!(it.next_chunk::<4>().unwrap_err().as_slice(), &[0, 1]);
+}
+
 #[test]
 fn test_collect_into_tuples() {
     let a = vec![(1, 2, 3), (4, 5, 6), (7, 8, 9)];