Skip to content

Commit

Permalink
Auto merge of #93700 - rossmacarthur:ft/iter-next-chunk, r=m-ou-se
Browse files Browse the repository at this point in the history
Add `Iterator::next_chunk`

See also rust-lang/rust#92393

### Prior art

-  [`Itertools::next_tuple()`](https://docs.rs/itertools/latest/itertools/trait.Itertools.html#method.next_tuple)

### Unresolved questions

- Should we also add `next_chunk_back` to `DoubleEndedIterator`?
- Should we rather call this `next_array()` or `next_array_chunk`?
  • Loading branch information
bors committed Jun 25, 2022
2 parents e357964 + cf655d6 commit 67c9bf3
Show file tree
Hide file tree
Showing 4 changed files with 100 additions and 31 deletions.
79 changes: 48 additions & 31 deletions core/src/array/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -780,24 +780,27 @@ where
}

/// Pulls `N` items from `iter` and returns them as an array. If the iterator
/// yields fewer than `N` items, `None` is returned and all already yielded
/// items are dropped.
/// yields fewer than `N` items, `Err` is returned containing an iterator over
/// the already yielded items.
///
/// Since the iterator is passed as a mutable reference and this function calls
/// `next` at most `N` times, the iterator can still be used afterwards to
/// retrieve the remaining items.
///
/// If `iter.next()` panicks, all items already yielded by the iterator are
/// dropped.
fn try_collect_into_array<I, T, R, const N: usize>(iter: &mut I) -> Option<R::TryType>
#[inline]
fn try_collect_into_array<I, T, R, const N: usize>(
iter: &mut I,
) -> Result<R::TryType, IntoIter<T, N>>
where
I: Iterator,
I::Item: Try<Output = T, Residual = R>,
R: Residual<[T; N]>,
{
if N == 0 {
// SAFETY: An empty array is always inhabited and has no validity invariants.
return unsafe { Some(Try::from_output(mem::zeroed())) };
return Ok(Try::from_output(unsafe { mem::zeroed() }));
}

struct Guard<'a, T, const N: usize> {
Expand All @@ -821,35 +824,49 @@ where
let mut array = MaybeUninit::uninit_array::<N>();
let mut guard = Guard { array_mut: &mut array, initialized: 0 };

while let Some(item_rslt) = iter.next() {
let item = match item_rslt.branch() {
ControlFlow::Break(r) => {
return Some(FromResidual::from_residual(r));
for _ in 0..N {
match iter.next() {
Some(item_rslt) => {
let item = match item_rslt.branch() {
ControlFlow::Break(r) => {
return Ok(FromResidual::from_residual(r));
}
ControlFlow::Continue(elem) => elem,
};

// SAFETY: `guard.initialized` starts at 0, is increased by one in the
// loop and the loop is aborted once it reaches N (which is
// `array.len()`).
unsafe {
guard.array_mut.get_unchecked_mut(guard.initialized).write(item);
}
guard.initialized += 1;
}
None => {
let alive = 0..guard.initialized;
mem::forget(guard);
// SAFETY: `array` was initialized with exactly `initialized`
// number of elements.
return Err(unsafe { IntoIter::new_unchecked(array, alive) });
}
ControlFlow::Continue(elem) => elem,
};

// SAFETY: `guard.initialized` starts at 0, is increased by one in the
// loop and the loop is aborted once it reaches N (which is
// `array.len()`).
unsafe {
guard.array_mut.get_unchecked_mut(guard.initialized).write(item);
}
guard.initialized += 1;

// Check if the whole array was initialized.
if guard.initialized == N {
mem::forget(guard);

// SAFETY: the condition above asserts that all elements are
// initialized.
let out = unsafe { MaybeUninit::array_assume_init(array) };
return Some(Try::from_output(out));
}
}

// This is only reached if the iterator is exhausted before
// `guard.initialized` reaches `N`. Also note that `guard` is dropped here,
// dropping all already initialized elements.
None
mem::forget(guard);
// SAFETY: All elements of the array were populated in the loop above.
let output = unsafe { MaybeUninit::array_assume_init(array) };
Ok(Try::from_output(output))
}

/// Returns the next chunk of `N` items from the iterator or errors with an
/// iterator over the remainder. Used for `Iterator::next_chunk`.
#[inline]
pub(crate) fn iter_next_chunk<I, const N: usize>(
iter: &mut I,
) -> Result<[I::Item; N], IntoIter<I::Item, N>>
where
I: Iterator,
{
let mut map = iter.map(NeverShortCircuit);
try_collect_into_array(&mut map).map(|NeverShortCircuit(arr)| arr)
}
42 changes: 42 additions & 0 deletions core/src/iter/traits/iterator.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use crate::array;
use crate::cmp::{self, Ordering};
use crate::ops::{ChangeOutputType, ControlFlow, FromResidual, Residual, Try};

Expand Down Expand Up @@ -102,6 +103,47 @@ pub trait Iterator {
#[stable(feature = "rust1", since = "1.0.0")]
fn next(&mut self) -> Option<Self::Item>;

/// Advances the iterator and returns an array containing the next `N` values.
///
/// If there are not enough elements to fill the array then `Err` is returned
/// containing an iterator over the remaining elements.
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// #![feature(iter_next_chunk)]
///
/// let mut iter = "lorem".chars();
///
/// assert_eq!(iter.next_chunk().unwrap(), ['l', 'o']); // N is inferred as 2
/// assert_eq!(iter.next_chunk().unwrap(), ['r', 'e', 'm']); // N is inferred as 3
/// assert_eq!(iter.next_chunk::<4>().unwrap_err().as_slice(), &[]); // N is explicitly 4
/// ```
///
/// Split a string and get the first three items.
///
/// ```
/// #![feature(iter_next_chunk)]
///
/// let quote = "not all those who wander are lost";
/// let [first, second, third] = quote.split_whitespace().next_chunk().unwrap();
/// assert_eq!(first, "not");
/// assert_eq!(second, "all");
/// assert_eq!(third, "those");
/// ```
#[inline]
#[unstable(feature = "iter_next_chunk", reason = "recently added", issue = "98326")]
fn next_chunk<const N: usize>(
&mut self,
) -> Result<[Self::Item; N], array::IntoIter<Self::Item, N>>
where
Self: Sized,
{
array::iter_next_chunk(self)
}

/// Returns the bounds on the remaining length of the iterator.
///
/// Specifically, `size_hint()` returns a tuple where the first element
Expand Down
9 changes: 9 additions & 0 deletions core/tests/iter/traits/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -575,6 +575,15 @@ fn iter_try_collect_uses_try_fold_not_next() {
// validation is just that it didn't panic.
}

#[test]
fn test_next_chunk() {
let mut it = 0..12;
assert_eq!(it.next_chunk().unwrap(), [0, 1, 2, 3]);
assert_eq!(it.next_chunk().unwrap(), []);
assert_eq!(it.next_chunk().unwrap(), [4, 5, 6, 7, 8, 9]);
assert_eq!(it.next_chunk::<4>().unwrap_err().as_slice(), &[10, 11]);
}

// just tests by whether or not this compiles
fn _empty_impl_all_auto_traits<T>() {
use std::panic::{RefUnwindSafe, UnwindSafe};
Expand Down
1 change: 1 addition & 0 deletions core/tests/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@
#![feature(iter_partition_in_place)]
#![feature(iter_intersperse)]
#![feature(iter_is_partitioned)]
#![feature(iter_next_chunk)]
#![feature(iter_order_by)]
#![feature(iterator_try_collect)]
#![feature(iterator_try_reduce)]
Expand Down

0 comments on commit 67c9bf3

Please sign in to comment.