Skip to content

Commit

Permalink
Rollup merge of rust-lang#100026 - WaffleLapkin:array-chunks, r=scottmcm
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
Dylan-DPC authored Aug 13, 2022
2 parents 6189cba + 5fbcde1 commit 2dac0f1
Show file tree
Hide file tree
Showing 7 changed files with 435 additions and 1 deletion.
182 changes: 182 additions & 0 deletions library/core/src/iter/adapters/array_chunks.rs
Original file line number Diff line number Diff line change
@@ -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
}
}
4 changes: 4 additions & 0 deletions library/core/src/iter/adapters/mod.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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;

Expand Down
2 changes: 2 additions & 0 deletions library/core/src/iter/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")]
Expand Down
45 changes: 44 additions & 1 deletion library/core/src/iter/traits/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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.
Expand Down
Loading

0 comments on commit 2dac0f1

Please sign in to comment.