Skip to content

Add Iterator::array_chunks method #87776

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 8 commits into from
Closed
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 51 additions & 0 deletions library/core/benches/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -367,3 +367,54 @@ fn bench_partial_cmp(b: &mut Bencher) {
fn bench_lt(b: &mut Bencher) {
b.iter(|| (0..100000).map(black_box).lt((0..100000).map(black_box)))
}

#[bench]
fn bench_array_map(b: &mut Bencher) {
b.iter(|| {
let mut acc = 0;
let iter = (0i64..10000).map(black_box).map(|_| black_box([0i64; 100]));
for_each_fold(iter, |x| acc += x.iter().sum::<i64>());
acc
});
}

#[bench]
fn bench_iter_array_chunks_loop(b: &mut Bencher) {
b.iter(|| {
let mut acc = 0;
let iter = (0i64..1000000).array_chunks::<100>().map(black_box);
for_each_loop(iter, |x| acc += x.iter().sum::<i64>());
acc
});
}

#[bench]
fn bench_iter_array_chunks_fold(b: &mut Bencher) {
b.iter(|| {
let mut acc = 0;
let iter = (0i64..1000000).array_chunks::<100>().map(black_box);
for_each_fold(iter, |x| acc += x.iter().sum::<i64>());
acc
});
}

#[bench]
fn bench_iter_array_chunks_ref_fold(b: &mut Bencher) {
b.iter(|| {
let mut acc = 0;
let mut iter = (0i64..1000000).array_chunks::<100>().map(black_box);
for_each_fold(iter.by_ref(), |x| acc += x.iter().sum::<i64>());
acc
});
}

#[bench]
fn bench_slice_array_chunks(b: &mut Bencher) {
let vec: Vec<_> = (0i64..1000000).collect();
b.iter(|| {
let mut acc = 0;
let iter = vec.array_chunks::<100>().map(black_box);
for_each_loop(iter, |x| acc += x.iter().sum::<i64>());
acc
});
}
2 changes: 2 additions & 0 deletions library/core/benches/lib.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
// wasm32 does not support benches (no time).
#![cfg(not(target_arch = "wasm32"))]
#![feature(array_chunks)]
#![feature(flt2dec)]
#![feature(iter_array_chunks)]
#![feature(test)]

extern crate test;
Expand Down
232 changes: 232 additions & 0 deletions library/core/src/iter/adapters/array_chunks.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,232 @@
use core::iter::FusedIterator;
use core::mem::{self, MaybeUninit};
use core::ops::{ControlFlow, Try};
use core::ptr;

#[derive(Debug)]
struct Buffer<T, const N: usize> {
array: [MaybeUninit<T>; N],
init: usize,
}

impl<T, const N: usize> Buffer<T, N> {
fn new() -> Self {
Self { array: MaybeUninit::uninit_array(), init: 0 }
}
}

impl<T: Clone, const N: usize> Clone for Buffer<T, N> {
fn clone(&self) -> Self {
let mut new = Self::new();
// SAFETY: this raw slice contains only the initialized objects.
let src = unsafe { MaybeUninit::slice_assume_init_ref(&self.array[..self.init]) };
MaybeUninit::write_slice_cloned(&mut new.array[..self.init], src);
new.init = self.init;
new
}
}

impl<T, const N: usize> Drop for Buffer<T, N> {
fn drop(&mut self) {
debug_assert!(self.init <= N);

let initialized_part = &mut self.array[..self.init];

// SAFETY: this raw slice will contain only the initialized objects
// that have not been dropped or moved.
unsafe {
ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(initialized_part));
}
}
}

// FIXME: Combine with `Guard` in `collect_into_array`.
struct Guard<T, const N: usize> {
ptr: *mut T,
init: usize,
}

impl<T, const N: usize> Drop for Guard<T, N> {
fn drop(&mut self) {
debug_assert!(self.init <= N);

let initialized_part = crate::ptr::slice_from_raw_parts_mut(self.ptr, self.init);

// SAFETY: this raw slice will contain only initialized objects.
unsafe {
crate::ptr::drop_in_place(initialized_part);
}
}
}

impl<T, const N: usize> Guard<T, N> {
fn new(array: &mut [MaybeUninit<T>; N]) -> Self {
Self { ptr: MaybeUninit::slice_as_mut_ptr(array), init: 0 }
}

fn with<R, F>(buffer: &mut Buffer<T, N>, f: F) -> R
where
F: FnOnce(&mut [MaybeUninit<T>; N], &mut usize) -> R,
{
let mut array = MaybeUninit::uninit_array();
let mut guard = Self::new(&mut array);
if buffer.init > 0 {
array = mem::replace(&mut buffer.array, MaybeUninit::uninit_array());
guard.init = mem::replace(&mut buffer.init, 0);
}
let res = f(&mut array, &mut guard.init);
if guard.init > 0 {
buffer.array = array;
buffer.init = guard.init;
mem::forget(guard);
}
res
}
}

/// An iterator that yields the elements of another iterator in
/// chunks of size `N`.
///
/// This `struct` is created by the [`array_chunks`] method on [`Iterator`]. See
/// its documentation for more.
///
/// [`array_chunks`]: Iterator::array_chunks
#[unstable(feature = "iter_array_chunks", issue = "none")]
#[derive(Debug, Clone)]
pub struct ArrayChunks<I: Iterator, const N: usize> {
iter: I,
buffer: Buffer<I::Item, N>,
}

impl<I: Iterator, const N: usize> ArrayChunks<I, N> {
pub(in crate::iter) fn new(iter: I) -> Self {
Self { iter, buffer: Buffer::new() }
}

/// Returns the remainder of the elements yielded by the original
/// iterator that were insufficient to fill another chunk. The
/// returned slice has at most `N-1` elements.
#[unstable(feature = "iter_array_chunks", issue = "none")]
pub fn remainder(&self) -> &[I::Item] {
// SAFETY: We know that all elements before `init` are properly initialized.
unsafe { MaybeUninit::slice_assume_init_ref(&self.buffer.array[..self.buffer.init]) }
}

/// Returns the remainder of the elements yielded by the original
/// iterator that were insufficient to fill another chunk. The
/// returned slice has at most `N-1` elements.
#[unstable(feature = "iter_array_chunks", issue = "none")]
pub fn remainder_mut(&mut self) -> &mut [I::Item] {
// SAFETY: We know that all elements before `init` are properly initialized.
unsafe { MaybeUninit::slice_assume_init_mut(&mut self.buffer.array[..self.buffer.init]) }
}
}

#[unstable(feature = "iter_array_chunks", issue = "none")]
impl<I: Iterator, const N: usize> Iterator for ArrayChunks<I, N> {
type Item = [I::Item; N];

fn next(&mut self) -> Option<Self::Item> {
let iter = &mut self.iter;
Guard::with(&mut self.buffer, |array, init| {
for slot in &mut array[*init..] {
slot.write(iter.next()?);
*init += 1;
}
*init = 0;
// SAFETY: The entire array has just been initialized.
unsafe {
Some(MaybeUninit::array_assume_init(mem::replace(
array,
MaybeUninit::uninit_array(),
)))
}
})
}

fn size_hint(&self) -> (usize, Option<usize>) {
let (lower, upper) = self.iter.size_hint();
(lower / N, upper.map(|x| x / N))
}

fn advance_by(&mut self, n: usize) -> Result<(), usize> {
let res = match n.checked_mul(N) {
Some(n) => self.iter.advance_by(n),
None => {
let n = (usize::MAX / N) * N;
self.iter.advance_by(n).and(Err(n))
}
};
res.map_err(|k| k / N)
}

fn try_fold<Acc, Fold, R>(&mut self, acc: Acc, mut fold: Fold) -> R
where
Self: Sized,
Fold: FnMut(Acc, Self::Item) -> R,
R: Try<Output = Acc>,
{
let iter = &mut self.iter;
Guard::with(&mut self.buffer, |array, init| {
let result = iter.try_fold(acc, |mut acc, x| {
// SAFETY: `init` starts at 0, is increased by one each iteration
// until it equals N (which is `array.len()`) and is reset to 0.
unsafe {
array.get_unchecked_mut(*init).write(x);
}
*init += 1;

if *init == N {
*init = 0;
// SAFETY: The entire array has just been initialized.
let array = unsafe {
MaybeUninit::array_assume_init(mem::replace(
array,
MaybeUninit::uninit_array(),
))
};
acc = fold(acc, array).branch()?;
}
ControlFlow::Continue(acc)
});

match result {
ControlFlow::Continue(acc) => R::from_output(acc),
ControlFlow::Break(res) => R::from_residual(res),
}
})
}

fn fold<Acc, Fold>(self, acc: Acc, mut fold: Fold) -> Acc
where
Fold: FnMut(Acc, Self::Item) -> Acc,
{
let Self { iter, mut buffer } = self;
Guard::with(&mut buffer, |array, init| {
iter.fold(acc, |mut acc, x| {
// SAFETY: `init` starts at 0, is increased by one each iteration
// until it equals N (which is `array.len()`) and is reset to 0.
unsafe {
array.get_unchecked_mut(*init).write(x);
}
*init += 1;

if *init == N {
*init = 0;
// SAFETY: The entire array has just been initialized.
let array = unsafe {
MaybeUninit::array_assume_init(mem::replace(
array,
MaybeUninit::uninit_array(),
))
};
acc = fold(acc, array);
}
acc
})
})
}
}

#[unstable(feature = "iter_array_chunks", issue = "none")]
impl<I: FusedIterator, const N: usize> FusedIterator for ArrayChunks<I, 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::{ControlFlow, Try};

mod array_chunks;
mod chain;
mod cloned;
mod copied;
Expand Down Expand Up @@ -30,6 +31,9 @@ pub use self::{
scan::Scan, skip::Skip, skip_while::SkipWhile, take::Take, take_while::TakeWhile, zip::Zip,
};

#[unstable(feature = "iter_array_chunks", issue = "none")]
pub use self::array_chunks::ArrayChunks;

#[stable(feature = "iter_cloned", since = "1.1.0")]
pub use self::cloned::Cloned;

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 @@ -393,6 +393,8 @@ pub use self::traits::{

#[unstable(feature = "iter_zip", issue = "83574")]
pub use self::adapters::zip;
#[unstable(feature = "iter_array_chunks", issue = "none")]
pub use self::adapters::ArrayChunks;
#[stable(feature = "iter_cloned", since = "1.1.0")]
pub use self::adapters::Cloned;
#[stable(feature = "iter_copied", since = "1.36.0")]
Expand Down
31 changes: 31 additions & 0 deletions library/core/src/iter/traits/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
use crate::cmp::{self, Ordering};
use crate::ops::{ControlFlow, Try};

use super::super::ArrayChunks;
use super::super::TrustedRandomAccessNoCoerce;
use super::super::{Chain, Cloned, Copied, Cycle, Enumerate, Filter, FilterMap, Fuse};
use super::super::{FlatMap, Flatten};
Expand Down Expand Up @@ -3468,6 +3469,36 @@ pub trait Iterator {
{
unreachable!("Always specialized");
}

/// Creates an iterator which yields arrays of `N` elements yielded by
/// the original iterator.
///
/// # Panics
///
/// Panics if `N` is zero.
///
/// # Examples
///
/// ```
/// #![feature(iter_array_chunks)]
/// let mut iter = (0..10).array_chunks::<3>();
///
/// assert_eq!(iter.next(), Some([0, 1, 2]));
/// assert_eq!(iter.next(), Some([3, 4, 5]));
/// assert_eq!(iter.next(), Some([6, 7, 8]));
/// assert_eq!(iter.next(), None);
///
/// assert_eq!(iter.remainder(), &[9]);
/// ```
#[inline]
#[unstable(feature = "iter_array_chunks", issue = "none")]
fn array_chunks<const N: usize>(self) -> ArrayChunks<Self, N>
where
Self: Sized,
{
assert_ne!(N, 0);
ArrayChunks::new(self)
}
}

#[stable(feature = "rust1", since = "1.0.0")]
Expand Down
Loading