diff --git a/library/core/src/iter/adapters/array_chunks.rs b/library/core/src/iter/adapters/array_chunks.rs new file mode 100644 index 0000000000000..9b479a9f8adfb --- /dev/null +++ b/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 { + iter: I, + remainder: Option>, +} + +impl ArrayChunks +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> { + self.remainder + } +} + +#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")] +impl Iterator for ArrayChunks +where + I: Iterator, +{ + type Item = [I::Item; N]; + + #[inline] + fn next(&mut self) -> Option { + self.try_for_each(ControlFlow::Break).break_value() + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + 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(&mut self, init: B, mut f: F) -> R + where + Self: Sized, + F: FnMut(B, Self::Item) -> R, + R: Try, + { + 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(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 DoubleEndedIterator for ArrayChunks +where + I: DoubleEndedIterator + ExactSizeIterator, +{ + #[inline] + fn next_back(&mut self) -> Option { + self.try_rfold((), |(), x| ControlFlow::Break(x)).break_value() + } + + fn try_rfold(&mut self, init: B, mut f: F) -> R + where + Self: Sized, + F: FnMut(B, Self::Item) -> R, + R: Try, + { + // 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(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 ArrayChunks +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 FusedIterator for ArrayChunks where I: FusedIterator {} + +#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")] +impl ExactSizeIterator for ArrayChunks +where + I: ExactSizeIterator, +{ + #[inline] + fn len(&self) -> usize { + self.iter.len() / N + } + + #[inline] + fn is_empty(&self) -> bool { + self.iter.len() < N + } +} diff --git a/library/core/src/iter/adapters/mod.rs b/library/core/src/iter/adapters/mod.rs index 916a26e242466..bf4fabad32a37 100644 --- a/library/core/src/iter/adapters/mod.rs +++ b/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; diff --git a/library/core/src/iter/mod.rs b/library/core/src/iter/mod.rs index d5c6aed5b6c8a..9514466bd0c05 100644 --- a/library/core/src/iter/mod.rs +++ b/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")] diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs index 275412b57b55f..b2d08f4b0f67b 100644 --- a/library/core/src/iter/traits/iterator.rs +++ b/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(self) -> ArrayChunks + where + Self: Sized, + { + ArrayChunks::new(self) + } + /// Sums the elements of an iterator. /// /// Takes each element, adds them together, and returns the result. diff --git a/library/core/tests/iter/adapters/array_chunks.rs b/library/core/tests/iter/adapters/array_chunks.rs new file mode 100644 index 0000000000000..4e9d89e1e580f --- /dev/null +++ b/library/core/tests/iter/adapters/array_chunks.rs @@ -0,0 +1,179 @@ +use core::cell::Cell; +use core::iter::{self, Iterator}; + +use super::*; + +#[test] +fn test_iterator_array_chunks_infer() { + let xs = [1, 1, 2, -2, 6, 0, 3, 1]; + for [a, b, c] in xs.iter().copied().array_chunks() { + assert_eq!(a + b + c, 4); + } +} + +#[test] +fn test_iterator_array_chunks_clone_and_drop() { + let count = Cell::new(0); + let mut it = (0..5).map(|_| CountDrop::new(&count)).array_chunks::<3>(); + assert_eq!(it.by_ref().count(), 1); + assert_eq!(count.get(), 3); + let mut it2 = it.clone(); + assert_eq!(count.get(), 3); + assert_eq!(it.into_remainder().unwrap().len(), 2); + assert_eq!(count.get(), 5); + assert!(it2.next().is_none()); + assert_eq!(it2.into_remainder().unwrap().len(), 2); + assert_eq!(count.get(), 7); +} + +#[test] +fn test_iterator_array_chunks_remainder() { + let mut it = (0..11).array_chunks::<4>(); + assert_eq!(it.next(), Some([0, 1, 2, 3])); + assert_eq!(it.next(), Some([4, 5, 6, 7])); + assert_eq!(it.next(), None); + assert_eq!(it.into_remainder().unwrap().as_slice(), &[8, 9, 10]); +} + +#[test] +fn test_iterator_array_chunks_size_hint() { + let it = (0..6).array_chunks::<1>(); + assert_eq!(it.size_hint(), (6, Some(6))); + + let it = (0..6).array_chunks::<3>(); + assert_eq!(it.size_hint(), (2, Some(2))); + + let it = (0..6).array_chunks::<5>(); + assert_eq!(it.size_hint(), (1, Some(1))); + + let it = (0..6).array_chunks::<7>(); + assert_eq!(it.size_hint(), (0, Some(0))); + + let it = (1..).array_chunks::<2>(); + assert_eq!(it.size_hint(), (usize::MAX / 2, None)); + + let it = (1..).filter(|x| x % 2 != 0).array_chunks::<2>(); + assert_eq!(it.size_hint(), (0, None)); +} + +#[test] +fn test_iterator_array_chunks_count() { + let it = (0..6).array_chunks::<1>(); + assert_eq!(it.count(), 6); + + let it = (0..6).array_chunks::<3>(); + assert_eq!(it.count(), 2); + + let it = (0..6).array_chunks::<5>(); + assert_eq!(it.count(), 1); + + let it = (0..6).array_chunks::<7>(); + assert_eq!(it.count(), 0); + + let it = (0..6).filter(|x| x % 2 == 0).array_chunks::<2>(); + assert_eq!(it.count(), 1); + + let it = iter::empty::().array_chunks::<2>(); + assert_eq!(it.count(), 0); + + let it = [(); usize::MAX].iter().array_chunks::<2>(); + assert_eq!(it.count(), usize::MAX / 2); +} + +#[test] +fn test_iterator_array_chunks_next_and_next_back() { + let mut it = (0..11).array_chunks::<3>(); + assert_eq!(it.next(), Some([0, 1, 2])); + assert_eq!(it.next_back(), Some([6, 7, 8])); + assert_eq!(it.next(), Some([3, 4, 5])); + assert_eq!(it.next_back(), None); + assert_eq!(it.next(), None); + assert_eq!(it.next_back(), None); + assert_eq!(it.next(), None); + assert_eq!(it.into_remainder().unwrap().as_slice(), &[9, 10]); +} + +#[test] +fn test_iterator_array_chunks_rev_remainder() { + let mut it = (0..11).array_chunks::<4>(); + { + let mut it = it.by_ref().rev(); + assert_eq!(it.next(), Some([4, 5, 6, 7])); + assert_eq!(it.next(), Some([0, 1, 2, 3])); + assert_eq!(it.next(), None); + assert_eq!(it.next(), None); + } + assert_eq!(it.into_remainder().unwrap().as_slice(), &[8, 9, 10]); +} + +#[test] +fn test_iterator_array_chunks_try_fold() { + let count = Cell::new(0); + let mut it = (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>(); + let result: Result<_, ()> = it.by_ref().try_fold(0, |acc, _item| Ok(acc + 1)); + assert_eq!(result, Ok(3)); + assert_eq!(count.get(), 9); + drop(it); + assert_eq!(count.get(), 10); + + let count = Cell::new(0); + let mut it = (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>(); + let result = it.by_ref().try_fold(0, |acc, _item| if acc < 2 { Ok(acc + 1) } else { Err(acc) }); + assert_eq!(result, Err(2)); + assert_eq!(count.get(), 9); + drop(it); + assert_eq!(count.get(), 9); +} + +#[test] +fn test_iterator_array_chunks_fold() { + let result = (1..11).array_chunks::<3>().fold(0, |acc, [a, b, c]| { + assert_eq!(acc + 1, a); + assert_eq!(acc + 2, b); + assert_eq!(acc + 3, c); + acc + 3 + }); + assert_eq!(result, 9); + + let count = Cell::new(0); + let result = + (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>().fold(0, |acc, _item| acc + 1); + assert_eq!(result, 3); + assert_eq!(count.get(), 10); +} + +#[test] +fn test_iterator_array_chunks_try_rfold() { + let count = Cell::new(0); + let mut it = (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>(); + let result: Result<_, ()> = it.try_rfold(0, |acc, _item| Ok(acc + 1)); + assert_eq!(result, Ok(3)); + assert_eq!(count.get(), 9); + drop(it); + assert_eq!(count.get(), 10); + + let count = Cell::new(0); + let mut it = (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>(); + let result = it.try_rfold(0, |acc, _item| if acc < 2 { Ok(acc + 1) } else { Err(acc) }); + assert_eq!(result, Err(2)); + assert_eq!(count.get(), 9); + drop(it); + assert_eq!(count.get(), 10); +} + +#[test] +fn test_iterator_array_chunks_rfold() { + let result = (1..11).array_chunks::<3>().rfold(0, |acc, [a, b, c]| { + assert_eq!(10 - (acc + 1), c); + assert_eq!(10 - (acc + 2), b); + assert_eq!(10 - (acc + 3), a); + acc + 3 + }); + assert_eq!(result, 9); + + let count = Cell::new(0); + let result = + (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>().rfold(0, |acc, _item| acc + 1); + assert_eq!(result, 3); + assert_eq!(count.get(), 10); +} diff --git a/library/core/tests/iter/adapters/mod.rs b/library/core/tests/iter/adapters/mod.rs index 567d9fe49cade..96539c0c394e2 100644 --- a/library/core/tests/iter/adapters/mod.rs +++ b/library/core/tests/iter/adapters/mod.rs @@ -1,3 +1,4 @@ +mod array_chunks; mod chain; mod cloned; mod copied; @@ -183,3 +184,25 @@ impl Clone for CountClone { ret } } + +#[derive(Debug, Clone)] +struct CountDrop<'a> { + dropped: bool, + count: &'a Cell, +} + +impl<'a> CountDrop<'a> { + pub fn new(count: &'a Cell) -> Self { + Self { dropped: false, count } + } +} + +impl Drop for CountDrop<'_> { + fn drop(&mut self) { + if self.dropped { + panic!("double drop"); + } + self.dropped = true; + self.count.set(self.count.get() + 1); + } +} diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index df9b1073a0994..09f1500f564cf 100644 --- a/library/core/tests/lib.rs +++ b/library/core/tests/lib.rs @@ -62,6 +62,7 @@ #![feature(slice_partition_dedup)] #![feature(int_log)] #![feature(iter_advance_by)] +#![feature(iter_array_chunks)] #![feature(iter_collect_into)] #![feature(iter_partition_in_place)] #![feature(iter_intersperse)]