From bbdff1fff49a6f816aa35d9b84c79b2603f0961b Mon Sep 17 00:00:00 2001 From: Ross MacArthur Date: Tue, 21 Jun 2022 08:57:02 +0200 Subject: [PATCH] Add `Iterator::next_chunk` --- library/core/src/array/mod.rs | 79 +++++++++++++--------- library/core/src/iter/traits/iterator.rs | 42 ++++++++++++ library/core/tests/iter/traits/iterator.rs | 9 +++ library/core/tests/lib.rs | 1 + 4 files changed, 100 insertions(+), 31 deletions(-) diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index 2ea4458bf64..c9823a136bc 100644 --- a/library/core/src/array/mod.rs +++ b/library/core/src/array/mod.rs @@ -780,8 +780,8 @@ unsafe fn try_collect_into_array_unchecked(iter: &mut I } /// 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 @@ -789,7 +789,10 @@ unsafe fn try_collect_into_array_unchecked(iter: &mut I /// /// If `iter.next()` panicks, all items already yielded by the iterator are /// dropped. -fn try_collect_into_array(iter: &mut I) -> Option +#[inline] +fn try_collect_into_array( + iter: &mut I, +) -> Result> where I: Iterator, I::Item: Try, @@ -797,7 +800,7 @@ fn try_collect_into_array(iter: &mut I) -> Option { @@ -821,35 +824,49 @@ fn drop(&mut self) { let mut array = MaybeUninit::uninit_array::(); 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( + iter: &mut I, +) -> Result<[I::Item; N], IntoIter> +where + I: Iterator, +{ + let mut map = iter.map(NeverShortCircuit); + try_collect_into_array(&mut map).map(|NeverShortCircuit(arr)| arr) } diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs index 1cc9133fc3d..326b98ec947 100644 --- a/library/core/src/iter/traits/iterator.rs +++ b/library/core/src/iter/traits/iterator.rs @@ -1,3 +1,4 @@ +use crate::array; use crate::cmp::{self, Ordering}; use crate::ops::{ChangeOutputType, ControlFlow, FromResidual, Residual, Try}; @@ -102,6 +103,47 @@ pub trait Iterator { #[stable(feature = "rust1", since = "1.0.0")] fn next(&mut self) -> Option; + /// 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( + &mut self, + ) -> Result<[Self::Item; N], array::IntoIter> + 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 diff --git a/library/core/tests/iter/traits/iterator.rs b/library/core/tests/iter/traits/iterator.rs index 731b1592d41..37345c1d381 100644 --- a/library/core/tests/iter/traits/iterator.rs +++ b/library/core/tests/iter/traits/iterator.rs @@ -575,6 +575,15 @@ fn try_fold(&mut self, init: B, f: F) -> R // 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() { use std::panic::{RefUnwindSafe, UnwindSafe}; diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index 63c9602abe7..9611e197a41 100644 --- a/library/core/tests/lib.rs +++ b/library/core/tests/lib.rs @@ -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)]