Rollup merge of rust-lang#100026 - WaffleLapkin:array-chunks, r=scottmcm

Add `Iterator::array_chunks` (take N+1)

A revival of rust-lang#92393.

r? `@Mark-Simulacrum`
cc `@rossmacarthur` `@scottmcm` `@the8472`

I've tried to address most of the review comments on the previous attempt. The only thing I didn't address is `try_fold` implementation, I've left the "custom" one for now, not sure what exactly should it use.

Author: Dylan-DPC

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

@@ -0,0 +1,182 @@
+use crate::array;
+use crate::iter::{ByRefSized, FusedIterator, Iterator};
+use crate::ops::{ControlFlow, NeverShortCircuit, Try};
+
+/// An iterator over `N` elements of the iterator at a time.
+///
+/// The chunks do not overlap. If `N` does not divide the length of the
+/// iterator, then the last up to `N-1` elements will be omitted.
+///
+/// This `struct` is created by the [`array_chunks`][Iterator::array_chunks]
+/// method on [`Iterator`]. See its documentation for more.
+#[derive(Debug, Clone)]
+#[must_use = "iterators are lazy and do nothing unless consumed"]
+#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
+pub struct ArrayChunks<I: Iterator, const N: usize> {
+    iter: I,
+    remainder: Option<array::IntoIter<I::Item, N>>,
+}
+
+impl<I, const N: usize> ArrayChunks<I, N>
+where
+    I: Iterator,
+{
+    #[track_caller]
+    pub(in crate::iter) fn new(iter: I) -> Self {
+        assert!(N != 0, "chunk size must be non-zero");
+        Self { iter, remainder: None }
+    }
+
+    /// Returns an iterator over the remaining elements of the original iterator
+    /// that are not going to be returned by this iterator. The returned
+    /// iterator will yield at most `N-1` elements.
+    #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
+    #[inline]
+    pub fn into_remainder(self) -> Option<array::IntoIter<I::Item, N>> {
+        self.remainder
+    }
+}
+
+#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
+impl<I, const N: usize> Iterator for ArrayChunks<I, N>
+where
+    I: Iterator,
+{
+    type Item = [I::Item; N];
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        self.try_for_each(ControlFlow::Break).break_value()
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let (lower, upper) = self.iter.size_hint();
+
+        (lower / N, upper.map(|n| n / N))
+    }
+
+    #[inline]
+    fn count(self) -> usize {
+        self.iter.count() / N
+    }
+
+    fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Output = B>,
+    {
+        let mut acc = init;
+        loop {
+            match self.iter.next_chunk() {
+                Ok(chunk) => acc = f(acc, chunk)?,
+                Err(remainder) => {
+                    // Make sure to not override `self.remainder` with an empty array
+                    // when `next` is called after `ArrayChunks` exhaustion.
+                    self.remainder.get_or_insert(remainder);
+
+                    break try { acc };
+                }
+            }
+        }
+    }
+
+    fn fold<B, F>(mut self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        self.try_fold(init, NeverShortCircuit::wrap_mut_2(f)).0
+    }
+}
+
+#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
+impl<I, const N: usize> DoubleEndedIterator for ArrayChunks<I, N>
+where
+    I: DoubleEndedIterator + ExactSizeIterator,
+{
+    #[inline]
+    fn next_back(&mut self) -> Option<Self::Item> {
+        self.try_rfold((), |(), x| ControlFlow::Break(x)).break_value()
+    }
+
+    fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Output = B>,
+    {
+        // We are iterating from the back we need to first handle the remainder.
+        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
+        // (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();
+            acc = f(acc, chunk)?
+        }
+
+        try { acc }
+    }
+
+    fn rfold<B, F>(mut self, init: B, f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        self.try_rfold(init, NeverShortCircuit::wrap_mut_2(f)).0
+    }
+}
+
+impl<I, const N: usize> ArrayChunks<I, N>
+where
+    I: DoubleEndedIterator + ExactSizeIterator,
+{
+    /// Updates `self.remainder` such that `self.iter.len` is divisible by `N`.
+    fn next_back_remainder(&mut self) {
+        // Make sure to not override `self.remainder` with an empty array
+        // when `next_back` is called after `ArrayChunks` exhaustion.
+        if self.remainder.is_some() {
+            return;
+        }
+
+        // We use the `ExactSizeIterator` implementation of the underlying
+        // iterator to know how many remaining elements there are.
+        let rem = self.iter.len() % N;
+
+        // Take the last `rem` elements out of `self.iter`.
+        let mut remainder =
+            // SAFETY: `unwrap_err` always succeeds because x % N < N for all x.
+            unsafe { self.iter.by_ref().rev().take(rem).next_chunk().unwrap_err_unchecked() };
+
+        // We used `.rev()` above, so we need to re-reverse the reminder
+        remainder.as_mut_slice().reverse();
+        self.remainder = Some(remainder);
+    }
+}
+
+#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
+impl<I, const N: usize> FusedIterator for ArrayChunks<I, N> where I: FusedIterator {}
+
+#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
+impl<I, const N: usize> ExactSizeIterator for ArrayChunks<I, N>
+where
+    I: ExactSizeIterator,
+{
+    #[inline]
+    fn len(&self) -> usize {
+        self.iter.len() / N
+    }
+
+    #[inline]
+    fn is_empty(&self) -> bool {
+        self.iter.len() < N
+    }
+}

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

@@ -1,6 +1,7 @@
 use crate::iter::{InPlaceIterable, Iterator};
 use crate::ops::{ChangeOutputType, ControlFlow, FromResidual, NeverShortCircuit, Residual, Try};
 
+mod array_chunks;
 mod by_ref_sized;
 mod chain;
 mod cloned;
@@ -32,6 +33,9 @@ pub use self::{
     scan::Scan, skip::Skip, skip_while::SkipWhile, take::Take, take_while::TakeWhile, zip::Zip,
 };
 
+#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
+pub use self::array_chunks::ArrayChunks;
+
 #[unstable(feature = "std_internals", issue = "none")]
 pub use self::by_ref_sized::ByRefSized;
 

library/core/src/iter/mod.rs

@@ -398,6 +398,8 @@ pub use self::traits::{
 
 #[stable(feature = "iter_zip", since = "1.59.0")]
 pub use self::adapters::zip;
+#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
+pub use self::adapters::ArrayChunks;
 #[unstable(feature = "std_internals", issue = "none")]
 pub use self::adapters::ByRefSized;
 #[stable(feature = "iter_cloned", since = "1.1.0")]

library/core/src/iter/traits/iterator.rs

@@ -5,7 +5,7 @@ use crate::ops::{ChangeOutputType, ControlFlow, FromResidual, Residual, Try};
 use super::super::try_process;
 use super::super::ByRefSized;
 use super::super::TrustedRandomAccessNoCoerce;
-use super::super::{Chain, Cloned, Copied, Cycle, Enumerate, Filter, FilterMap, Fuse};
+use super::super::{ArrayChunks, Chain, Cloned, Copied, Cycle, Enumerate, Filter, FilterMap, Fuse};
 use super::super::{FlatMap, Flatten};
 use super::super::{FromIterator, Intersperse, IntersperseWith, Product, Sum, Zip};
 use super::super::{
@@ -3316,6 +3316,49 @@ pub trait Iterator {
         Cycle::new(self)
     }
 
+    /// Returns an iterator over `N` elements of the iterator at a time.
+    ///
+    /// The chunks do not overlap. If `N` does not divide the length of the
+    /// iterator, then the last up to `N-1` elements will be omitted and can be
+    /// retrieved from the [`.into_remainder()`][ArrayChunks::into_remainder]
+    /// function of the iterator.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `N` is 0.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// #![feature(iter_array_chunks)]
+    ///
+    /// let mut iter = "lorem".chars().array_chunks();
+    /// assert_eq!(iter.next(), Some(['l', 'o']));
+    /// assert_eq!(iter.next(), Some(['r', 'e']));
+    /// assert_eq!(iter.next(), None);
+    /// assert_eq!(iter.into_remainder().unwrap().as_slice(), &['m']);
+    /// ```
+    ///
+    /// ```
+    /// #![feature(iter_array_chunks)]
+    ///
+    /// let data = [1, 1, 2, -2, 6, 0, 3, 1];
+    /// //          ^-----^  ^------^
+    /// for [x, y, z] in data.iter().array_chunks() {
+    ///     assert_eq!(x + y + z, 4);
+    /// }
+    /// ```
+    #[track_caller]
+    #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
+    fn array_chunks<const N: usize>(self) -> ArrayChunks<Self, N>
+    where
+        Self: Sized,
+    {
+        ArrayChunks::new(self)
+    }
+
     /// Sums the elements of an iterator.
     ///
     /// Takes each element, adds them together, and returns the result.