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

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

A revival of https://github.com/rust-lang/rust/pull/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.

7 files changed, 435 insertions, 1 deletions

diff --git a/library/core/src/iter/adapters/array_chunks.rs b/library/core/src/iter/adapters/array_chunks.rs
new file mode 100644
index 00000000000..9b479a9f8ad
--- /dev/null
+++ b/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
+    }
+}
diff --git a/library/core/src/iter/adapters/mod.rs b/library/core/src/iter/adapters/mod.rs
index 916a26e2424..bf4fabad32a 100644
--- a/library/core/src/iter/adapters/mod.rs
+++ b/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;
 
diff --git a/library/core/src/iter/mod.rs b/library/core/src/iter/mod.rs
index d5c6aed5b6c..9514466bd0c 100644
--- a/library/core/src/iter/mod.rs
+++ b/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")]
diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs
index 275412b57b5..b2d08f4b0f6 100644
--- a/library/core/src/iter/traits/iterator.rs
+++ b/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.
diff --git a/library/core/tests/iter/adapters/array_chunks.rs b/library/core/tests/iter/adapters/array_chunks.rs
new file mode 100644
index 00000000000..4e9d89e1e58
--- /dev/null
+++ b/library/core/tests/iter/adapters/array_chunks.rs
@@ -0,0 +1,179 @@
+use core::cell::Cell;
+use core::iter::{self, Iterator};
+
+use super::*;
+
+#[test]
+fn test_iterator_array_chunks_infer() {
+    let xs = [1, 1, 2, -2, 6, 0, 3, 1];
+    for [a, b, c] in xs.iter().copied().array_chunks() {
+        assert_eq!(a + b + c, 4);
+    }
+}
+
+#[test]
+fn test_iterator_array_chunks_clone_and_drop() {
+    let count = Cell::new(0);
+    let mut it = (0..5).map(|_| CountDrop::new(&count)).array_chunks::<3>();
+    assert_eq!(it.by_ref().count(), 1);
+    assert_eq!(count.get(), 3);
+    let mut it2 = it.clone();
+    assert_eq!(count.get(), 3);
+    assert_eq!(it.into_remainder().unwrap().len(), 2);
+    assert_eq!(count.get(), 5);
+    assert!(it2.next().is_none());
+    assert_eq!(it2.into_remainder().unwrap().len(), 2);
+    assert_eq!(count.get(), 7);
+}
+
+#[test]
+fn test_iterator_array_chunks_remainder() {
+    let mut it = (0..11).array_chunks::<4>();
+    assert_eq!(it.next(), Some([0, 1, 2, 3]));
+    assert_eq!(it.next(), Some([4, 5, 6, 7]));
+    assert_eq!(it.next(), None);
+    assert_eq!(it.into_remainder().unwrap().as_slice(), &[8, 9, 10]);
+}
+
+#[test]
+fn test_iterator_array_chunks_size_hint() {
+    let it = (0..6).array_chunks::<1>();
+    assert_eq!(it.size_hint(), (6, Some(6)));
+
+    let it = (0..6).array_chunks::<3>();
+    assert_eq!(it.size_hint(), (2, Some(2)));
+
+    let it = (0..6).array_chunks::<5>();
+    assert_eq!(it.size_hint(), (1, Some(1)));
+
+    let it = (0..6).array_chunks::<7>();
+    assert_eq!(it.size_hint(), (0, Some(0)));
+
+    let it = (1..).array_chunks::<2>();
+    assert_eq!(it.size_hint(), (usize::MAX / 2, None));
+
+    let it = (1..).filter(|x| x % 2 != 0).array_chunks::<2>();
+    assert_eq!(it.size_hint(), (0, None));
+}
+
+#[test]
+fn test_iterator_array_chunks_count() {
+    let it = (0..6).array_chunks::<1>();
+    assert_eq!(it.count(), 6);
+
+    let it = (0..6).array_chunks::<3>();
+    assert_eq!(it.count(), 2);
+
+    let it = (0..6).array_chunks::<5>();
+    assert_eq!(it.count(), 1);
+
+    let it = (0..6).array_chunks::<7>();
+    assert_eq!(it.count(), 0);
+
+    let it = (0..6).filter(|x| x % 2 == 0).array_chunks::<2>();
+    assert_eq!(it.count(), 1);
+
+    let it = iter::empty::<i32>().array_chunks::<2>();
+    assert_eq!(it.count(), 0);
+
+    let it = [(); usize::MAX].iter().array_chunks::<2>();
+    assert_eq!(it.count(), usize::MAX / 2);
+}
+
+#[test]
+fn test_iterator_array_chunks_next_and_next_back() {
+    let mut it = (0..11).array_chunks::<3>();
+    assert_eq!(it.next(), Some([0, 1, 2]));
+    assert_eq!(it.next_back(), Some([6, 7, 8]));
+    assert_eq!(it.next(), Some([3, 4, 5]));
+    assert_eq!(it.next_back(), None);
+    assert_eq!(it.next(), None);
+    assert_eq!(it.next_back(), None);
+    assert_eq!(it.next(), None);
+    assert_eq!(it.into_remainder().unwrap().as_slice(), &[9, 10]);
+}
+
+#[test]
+fn test_iterator_array_chunks_rev_remainder() {
+    let mut it = (0..11).array_chunks::<4>();
+    {
+        let mut it = it.by_ref().rev();
+        assert_eq!(it.next(), Some([4, 5, 6, 7]));
+        assert_eq!(it.next(), Some([0, 1, 2, 3]));
+        assert_eq!(it.next(), None);
+        assert_eq!(it.next(), None);
+    }
+    assert_eq!(it.into_remainder().unwrap().as_slice(), &[8, 9, 10]);
+}
+
+#[test]
+fn test_iterator_array_chunks_try_fold() {
+    let count = Cell::new(0);
+    let mut it = (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>();
+    let result: Result<_, ()> = it.by_ref().try_fold(0, |acc, _item| Ok(acc + 1));
+    assert_eq!(result, Ok(3));
+    assert_eq!(count.get(), 9);
+    drop(it);
+    assert_eq!(count.get(), 10);
+
+    let count = Cell::new(0);
+    let mut it = (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>();
+    let result = it.by_ref().try_fold(0, |acc, _item| if acc < 2 { Ok(acc + 1) } else { Err(acc) });
+    assert_eq!(result, Err(2));
+    assert_eq!(count.get(), 9);
+    drop(it);
+    assert_eq!(count.get(), 9);
+}
+
+#[test]
+fn test_iterator_array_chunks_fold() {
+    let result = (1..11).array_chunks::<3>().fold(0, |acc, [a, b, c]| {
+        assert_eq!(acc + 1, a);
+        assert_eq!(acc + 2, b);
+        assert_eq!(acc + 3, c);
+        acc + 3
+    });
+    assert_eq!(result, 9);
+
+    let count = Cell::new(0);
+    let result =
+        (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>().fold(0, |acc, _item| acc + 1);
+    assert_eq!(result, 3);
+    assert_eq!(count.get(), 10);
+}
+
+#[test]
+fn test_iterator_array_chunks_try_rfold() {
+    let count = Cell::new(0);
+    let mut it = (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>();
+    let result: Result<_, ()> = it.try_rfold(0, |acc, _item| Ok(acc + 1));
+    assert_eq!(result, Ok(3));
+    assert_eq!(count.get(), 9);
+    drop(it);
+    assert_eq!(count.get(), 10);
+
+    let count = Cell::new(0);
+    let mut it = (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>();
+    let result = it.try_rfold(0, |acc, _item| if acc < 2 { Ok(acc + 1) } else { Err(acc) });
+    assert_eq!(result, Err(2));
+    assert_eq!(count.get(), 9);
+    drop(it);
+    assert_eq!(count.get(), 10);
+}
+
+#[test]
+fn test_iterator_array_chunks_rfold() {
+    let result = (1..11).array_chunks::<3>().rfold(0, |acc, [a, b, c]| {
+        assert_eq!(10 - (acc + 1), c);
+        assert_eq!(10 - (acc + 2), b);
+        assert_eq!(10 - (acc + 3), a);
+        acc + 3
+    });
+    assert_eq!(result, 9);
+
+    let count = Cell::new(0);
+    let result =
+        (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>().rfold(0, |acc, _item| acc + 1);
+    assert_eq!(result, 3);
+    assert_eq!(count.get(), 10);
+}
diff --git a/library/core/tests/iter/adapters/mod.rs b/library/core/tests/iter/adapters/mod.rs
index 567d9fe49ca..96539c0c394 100644
--- a/library/core/tests/iter/adapters/mod.rs
+++ b/library/core/tests/iter/adapters/mod.rs
@@ -1,3 +1,4 @@
+mod array_chunks;
 mod chain;
 mod cloned;
 mod copied;
@@ -183,3 +184,25 @@ impl Clone for CountClone {
         ret
     }
 }
+
+#[derive(Debug, Clone)]
+struct CountDrop<'a> {
+    dropped: bool,
+    count: &'a Cell<usize>,
+}
+
+impl<'a> CountDrop<'a> {
+    pub fn new(count: &'a Cell<usize>) -> Self {
+        Self { dropped: false, count }
+    }
+}
+
+impl Drop for CountDrop<'_> {
+    fn drop(&mut self) {
+        if self.dropped {
+            panic!("double drop");
+        }
+        self.dropped = true;
+        self.count.set(self.count.get() + 1);
+    }
+}
diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs
index df9b1073a09..09f1500f564 100644
--- a/library/core/tests/lib.rs
+++ b/library/core/tests/lib.rs
@@ -62,6 +62,7 @@
 #![feature(slice_partition_dedup)]
 #![feature(int_log)]
 #![feature(iter_advance_by)]
+#![feature(iter_array_chunks)]
 #![feature(iter_collect_into)]
 #![feature(iter_partition_in_place)]
 #![feature(iter_intersperse)]