Add Iterator::array_chunks (take N+1)

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.