From 1f475477d573c2bef9848507ce5dfefa5ed625c8 Mon Sep 17 00:00:00 2001 From: Ross MacArthur Date: Sun, 2 Jan 2022 18:31:56 +0200 Subject: [PATCH 01/12] Add `Iterator::array_chunks()` --- core/src/iter/adapters/array_chunks.rs | 427 +++++++++++++++++++++++ core/src/iter/adapters/mod.rs | 4 + core/src/iter/mod.rs | 2 + core/src/iter/traits/iterator.rs | 42 ++- core/tests/iter/adapters/array_chunks.rs | 198 +++++++++++ core/tests/iter/adapters/mod.rs | 23 ++ core/tests/lib.rs | 1 + 7 files changed, 696 insertions(+), 1 deletion(-) create mode 100644 core/src/iter/adapters/array_chunks.rs create mode 100644 core/tests/iter/adapters/array_chunks.rs diff --git a/core/src/iter/adapters/array_chunks.rs b/core/src/iter/adapters/array_chunks.rs new file mode 100644 index 000000000..f9c3f03cb --- /dev/null +++ b/core/src/iter/adapters/array_chunks.rs @@ -0,0 +1,427 @@ +use crate::iter::{Fuse, FusedIterator, Iterator, TrustedLen}; +use crate::mem; +use crate::mem::MaybeUninit; +use crate::ops::{ControlFlow, Try}; +use crate::ptr; + +#[derive(Debug)] +struct Remainder { + array: [MaybeUninit; N], + init: usize, +} + +impl Remainder { + fn new() -> Self { + Self { array: MaybeUninit::uninit_array(), init: 0 } + } + + unsafe fn with_init(array: [MaybeUninit; N], init: usize) -> Self { + Self { array, init } + } + + fn as_slice(&self) -> &[T] { + debug_assert!(self.init <= N); + // SAFETY: This raw slice will only contain the initialized objects + // within the buffer. + unsafe { + let slice = self.array.get_unchecked(..self.init); + MaybeUninit::slice_assume_init_ref(slice) + } + } + + fn as_mut_slice(&mut self) -> &mut [T] { + debug_assert!(self.init <= N); + // SAFETY: This raw slice will only contain the initialized objects + // within the buffer. + unsafe { + let slice = self.array.get_unchecked_mut(..self.init); + MaybeUninit::slice_assume_init_mut(slice) + } + } +} + +impl Clone for Remainder +where + T: Clone, +{ + fn clone(&self) -> Self { + let mut new = Self::new(); + // SAFETY: The new array is the same size and `init` is always less than + // or equal to `N`. + let this = unsafe { new.array.get_unchecked_mut(..self.init) }; + MaybeUninit::write_slice_cloned(this, self.as_slice()); + new.init = self.init; + new + } +} + +impl Drop for Remainder { + fn drop(&mut self) { + // SAFETY: This raw slice will only contain the initialized objects + // within the buffer. + unsafe { ptr::drop_in_place(self.as_mut_slice()) } + } +} + +/// 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 = "none")] +pub struct ArrayChunks { + iter: Fuse, + remainder: Remainder, +} + +impl ArrayChunks +where + I: Iterator, +{ + pub(in crate::iter) fn new(iter: I) -> Self { + assert!(N != 0, "chunk size must be non-zero"); + Self { iter: iter.fuse(), remainder: Remainder::new() } + } + + /// Returns a reference to the remaining elements of the original iterator + /// that are not going to be returned by this iterator. The returned slice + /// has at most `N-1` elements. + #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] + #[inline] + pub fn remainder(&self) -> &[I::Item] { + self.remainder.as_slice() + } + + /// Returns a mutable reference to the remaining elements of the original + /// iterator that are not going to be returned by this iterator. The + /// returned slice has at most `N-1` elements. + #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] + #[inline] + pub fn remainder_mut(&mut self) -> &mut [I::Item] { + self.remainder.as_mut_slice() + } +} + +#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] +impl Iterator for ArrayChunks +where + I: Iterator, +{ + type Item = [I::Item; N]; + + #[inline] + fn next(&mut self) -> Option { + let mut array = MaybeUninit::uninit_array(); + // SAFETY: `array` will still be valid if `guard` is dropped. + let mut guard = unsafe { FrontGuard::new(&mut array) }; + + for slot in array.iter_mut() { + match self.iter.next() { + Some(item) => { + slot.write(item); + guard.init += 1; + } + None => { + if guard.init > 0 { + let init = guard.init; + mem::forget(guard); + // SAFETY: `array` was initialized with `init` elements. + self.remainder = unsafe { Remainder::with_init(array, init) }; + } + return None; + } + } + } + + mem::forget(guard); + // SAFETY: All elements of the array were populated in the loop above. + Some(unsafe { MaybeUninit::array_assume_init(array) }) + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + let (lower, upper) = self.iter.size_hint(); + // Keep infinite iterator size hint lower bound as `usize::MAX`. This + // is required to implement `TrustedLen`. + if lower == usize::MAX { + return (lower, upper); + } + (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 array = MaybeUninit::uninit_array(); + // SAFETY: `array` will still be valid if `guard` is dropped. + let mut guard = unsafe { FrontGuard::new(&mut array) }; + + let result = self.iter.try_fold(init, |mut acc, item| { + // SAFETY: `init` starts at 0, increases by one each iteration and + // is reset to 0 once it reaches N. + unsafe { array.get_unchecked_mut(guard.init) }.write(item); + guard.init += 1; + if guard.init == N { + guard.init = 0; + let array = mem::replace(&mut array, MaybeUninit::uninit_array()); + // SAFETY: the condition above asserts that all elements are + // initialized. + let item = unsafe { MaybeUninit::array_assume_init(array) }; + acc = f(acc, item)?; + } + R::from_output(acc) + }); + match result.branch() { + ControlFlow::Continue(o) => { + if guard.init > 0 { + let init = guard.init; + mem::forget(guard); + // SAFETY: `array` was initialized with `init` elements. + self.remainder = unsafe { Remainder::with_init(array, init) }; + } + R::from_output(o) + } + ControlFlow::Break(r) => R::from_residual(r), + } + } + + fn fold(self, init: B, mut f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B, + { + let mut array = MaybeUninit::uninit_array(); + // SAFETY: `array` will still be valid if `guard` is dropped. + let mut guard = unsafe { FrontGuard::new(&mut array) }; + + self.iter.fold(init, |mut acc, item| { + // SAFETY: `init` starts at 0, increases by one each iteration and + // is reset to 0 once it reaches N. + unsafe { array.get_unchecked_mut(guard.init) }.write(item); + guard.init += 1; + if guard.init == N { + guard.init = 0; + let array = mem::replace(&mut array, MaybeUninit::uninit_array()); + // SAFETY: the condition above asserts that all elements are + // initialized. + let item = unsafe { MaybeUninit::array_assume_init(array) }; + acc = f(acc, item); + } + acc + }) + } +} + +/// A guard for an array where elements are filled from the left. +struct FrontGuard { + /// A pointer to the array that is being filled. We need to use a raw + /// pointer here because of the lifetime issues in the fold implementations. + ptr: *mut T, + /// The number of *initialized* elements. + init: usize, +} + +impl FrontGuard { + unsafe fn new(array: &mut [MaybeUninit; N]) -> Self { + Self { ptr: MaybeUninit::slice_as_mut_ptr(array), init: 0 } + } +} + +impl Drop for FrontGuard { + fn drop(&mut self) { + debug_assert!(self.init <= N); + // SAFETY: This raw slice will only contain the initialized objects + // within the buffer. + unsafe { + let slice = ptr::slice_from_raw_parts_mut(self.ptr, self.init); + ptr::drop_in_place(slice); + } + } +} + +#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] +impl DoubleEndedIterator for ArrayChunks +where + I: DoubleEndedIterator + ExactSizeIterator, +{ + #[inline] + fn next_back(&mut self) -> Option { + // We are iterating from the back we need to first handle the remainder. + self.next_back_remainder()?; + + let mut array = MaybeUninit::uninit_array(); + // SAFETY: `array` will still be valid if `guard` is dropped. + let mut guard = unsafe { BackGuard::new(&mut array) }; + + for slot in array.iter_mut().rev() { + slot.write(self.iter.next_back()?); + guard.uninit -= 1; + } + + mem::forget(guard); + // SAFETY: All elements of the array were populated in the loop above. + Some(unsafe { MaybeUninit::array_assume_init(array) }) + } + + 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. + if self.next_back_remainder().is_none() { + return R::from_output(init); + } + + let mut array = MaybeUninit::uninit_array(); + // SAFETY: `array` will still be valid if `guard` is dropped. + let mut guard = unsafe { BackGuard::new(&mut array) }; + + self.iter.try_rfold(init, |mut acc, item| { + guard.uninit -= 1; + // SAFETY: `uninit` starts at N, decreases by one each iteration and + // is reset to N once it reaches 0. + unsafe { array.get_unchecked_mut(guard.uninit) }.write(item); + if guard.uninit == 0 { + guard.uninit = N; + let array = mem::replace(&mut array, MaybeUninit::uninit_array()); + // SAFETY: the condition above asserts that all elements are + // initialized. + let item = unsafe { MaybeUninit::array_assume_init(array) }; + acc = f(acc, item)?; + } + R::from_output(acc) + }) + } + + fn rfold(mut self, init: B, mut f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B, + { + // We are iterating from the back we need to first handle the remainder. + if self.next_back_remainder().is_none() { + return init; + } + + let mut array = MaybeUninit::uninit_array(); + + // SAFETY: `array` will still be valid if `guard` is dropped. + let mut guard = unsafe { BackGuard::new(&mut array) }; + + self.iter.rfold(init, |mut acc, item| { + guard.uninit -= 1; + // SAFETY: `uninit` starts at N, decreases by one each iteration and + // is reset to N once it reaches 0. + unsafe { array.get_unchecked_mut(guard.uninit) }.write(item); + if guard.uninit == 0 { + guard.uninit = N; + let array = mem::replace(&mut array, MaybeUninit::uninit_array()); + // SAFETY: the condition above asserts that all elements are + // initialized. + let item = unsafe { MaybeUninit::array_assume_init(array) }; + acc = f(acc, item); + } + acc + }) + } +} + +impl ArrayChunks +where + I: DoubleEndedIterator + ExactSizeIterator, +{ + #[inline] + fn next_back_remainder(&mut self) -> Option<()> { + // We use the `ExactSizeIterator` implementation of the underlying + // iterator to know how many remaining elements there are. + let rem = self.iter.len() % N; + if rem == 0 { + return Some(()); + } + + let mut array = MaybeUninit::uninit_array(); + + // SAFETY: The array will still be valid if `guard` is dropped and + // it is forgotten otherwise. + let mut guard = unsafe { FrontGuard::new(&mut array) }; + + // SAFETY: `rem` is in the range 1..N based on how it is calculated. + for slot in unsafe { array.get_unchecked_mut(..rem) }.iter_mut() { + slot.write(self.iter.next_back()?); + guard.init += 1; + } + + let init = guard.init; + mem::forget(guard); + // SAFETY: `array` was initialized with exactly `init` elements. + self.remainder = unsafe { + array.get_unchecked_mut(..init).reverse(); + Remainder::with_init(array, init) + }; + Some(()) + } +} + +/// A guard for an array where elements are filled from the right. +struct BackGuard { + /// A pointer to the array that is being filled. We need to use a raw + /// pointer here because of the lifetime issues in the rfold implementations. + ptr: *mut T, + /// The number of *uninitialized* elements. + uninit: usize, +} + +impl BackGuard { + unsafe fn new(array: &mut [MaybeUninit; N]) -> Self { + Self { ptr: MaybeUninit::slice_as_mut_ptr(array), uninit: N } + } +} + +impl Drop for BackGuard { + fn drop(&mut self) { + debug_assert!(self.uninit <= N); + // SAFETY: This raw slice will only contain the initialized objects + // within the buffer. + unsafe { + let ptr = self.ptr.offset(self.uninit as isize); + let slice = ptr::slice_from_raw_parts_mut(ptr, N - self.uninit); + ptr::drop_in_place(slice); + } + } +} + +#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] +impl FusedIterator for ArrayChunks where I: FusedIterator {} + +#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] +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 == 0 + } +} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl TrustedLen for ArrayChunks where I: TrustedLen {} diff --git a/core/src/iter/adapters/mod.rs b/core/src/iter/adapters/mod.rs index 916a26e24..39e7ab878 100644 --- a/core/src/iter/adapters/mod.rs +++ b/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 = "none")] +pub use self::array_chunks::ArrayChunks; + #[unstable(feature = "std_internals", issue = "none")] pub use self::by_ref_sized::ByRefSized; diff --git a/core/src/iter/mod.rs b/core/src/iter/mod.rs index d5c6aed5b..d48e3a52c 100644 --- a/core/src/iter/mod.rs +++ b/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 = "none")] +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/core/src/iter/traits/iterator.rs b/core/src/iter/traits/iterator.rs index 275412b57..8bf41ca6f 100644 --- a/core/src/iter/traits/iterator.rs +++ b/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,46 @@ 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. + /// + /// # 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.remainder(), &['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); + /// } + /// ``` + #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] + 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/core/tests/iter/adapters/array_chunks.rs b/core/tests/iter/adapters/array_chunks.rs new file mode 100644 index 000000000..6845c94d3 --- /dev/null +++ b/core/tests/iter/adapters/array_chunks.rs @@ -0,0 +1,198 @@ +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); + assert_eq!(it.remainder().len(), 2); + + let mut it2 = it.clone(); + assert_eq!(count.get(), 3); + assert_eq!(it2.remainder().len(), 2); + + drop(it); + assert_eq!(count.get(), 5); + assert_eq!(it2.remainder().len(), 2); + assert!(it2.next().is_none()); + + drop(it2); + assert_eq!(count.get(), 7); +} + +#[test] +fn test_iterator_array_chunks_remainder() { + let mut it = (0..11).array_chunks::<4>(); + assert_eq!(it.remainder(), &[]); + assert_eq!(it.remainder_mut(), &[]); + assert_eq!(it.next(), Some([0, 1, 2, 3])); + assert_eq!(it.remainder(), &[]); + assert_eq!(it.remainder_mut(), &[]); + assert_eq!(it.next(), Some([4, 5, 6, 7])); + assert_eq!(it.remainder(), &[]); + assert_eq!(it.remainder_mut(), &[]); + assert_eq!(it.next(), None); + assert_eq!(it.next(), None); + assert_eq!(it.remainder(), &[8, 9, 10]); + assert_eq!(it.remainder_mut(), &[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, 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.remainder(), &[9, 10]); + assert_eq!(it.remainder_mut(), &[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.remainder(), &[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!(it.remainder().len(), 1); + 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!(it.remainder().len(), 0); + 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!(it.remainder().len(), 1); + 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/core/tests/iter/adapters/mod.rs b/core/tests/iter/adapters/mod.rs index 567d9fe49..96539c0c3 100644 --- a/core/tests/iter/adapters/mod.rs +++ b/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/core/tests/lib.rs b/core/tests/lib.rs index db94368f6..8b4838bb7 100644 --- a/core/tests/lib.rs +++ b/core/tests/lib.rs @@ -61,6 +61,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)] From d5436e4ec928f19cc5049b3f2511f7e8b4b949fb Mon Sep 17 00:00:00 2001 From: Ross MacArthur Date: Fri, 4 Feb 2022 15:57:58 +0200 Subject: [PATCH 02/12] Use `array::IntoIter` for the `ArrayChunks` remainder --- core/src/array/iter.rs | 10 +++ core/src/iter/adapters/array_chunks.rs | 93 ++++-------------------- core/src/iter/traits/iterator.rs | 6 +- core/tests/iter/adapters/array_chunks.rs | 29 ++------ 4 files changed, 33 insertions(+), 105 deletions(-) diff --git a/core/src/array/iter.rs b/core/src/array/iter.rs index f4885ed9f..459cd094c 100644 --- a/core/src/array/iter.rs +++ b/core/src/array/iter.rs @@ -84,6 +84,16 @@ impl IntoIter { IntoIterator::into_iter(array) } + /// Creates a new iterator from a partially initalized array. + /// + /// # Safety + /// + /// The caller must guarantee that all and only the `alive` elements of + /// `data` are initialized. + pub(crate) unsafe fn with_partial(data: [MaybeUninit; N], alive: Range) -> Self { + Self { data, alive } + } + /// Creates an iterator over the elements in a partially-initialized buffer. /// /// If you have a fully-initialized array, then use [`IntoIterator`]. diff --git a/core/src/iter/adapters/array_chunks.rs b/core/src/iter/adapters/array_chunks.rs index f9c3f03cb..2ec1284c3 100644 --- a/core/src/iter/adapters/array_chunks.rs +++ b/core/src/iter/adapters/array_chunks.rs @@ -1,68 +1,10 @@ +use crate::array; use crate::iter::{Fuse, FusedIterator, Iterator, TrustedLen}; use crate::mem; use crate::mem::MaybeUninit; use crate::ops::{ControlFlow, Try}; use crate::ptr; -#[derive(Debug)] -struct Remainder { - array: [MaybeUninit; N], - init: usize, -} - -impl Remainder { - fn new() -> Self { - Self { array: MaybeUninit::uninit_array(), init: 0 } - } - - unsafe fn with_init(array: [MaybeUninit; N], init: usize) -> Self { - Self { array, init } - } - - fn as_slice(&self) -> &[T] { - debug_assert!(self.init <= N); - // SAFETY: This raw slice will only contain the initialized objects - // within the buffer. - unsafe { - let slice = self.array.get_unchecked(..self.init); - MaybeUninit::slice_assume_init_ref(slice) - } - } - - fn as_mut_slice(&mut self) -> &mut [T] { - debug_assert!(self.init <= N); - // SAFETY: This raw slice will only contain the initialized objects - // within the buffer. - unsafe { - let slice = self.array.get_unchecked_mut(..self.init); - MaybeUninit::slice_assume_init_mut(slice) - } - } -} - -impl Clone for Remainder -where - T: Clone, -{ - fn clone(&self) -> Self { - let mut new = Self::new(); - // SAFETY: The new array is the same size and `init` is always less than - // or equal to `N`. - let this = unsafe { new.array.get_unchecked_mut(..self.init) }; - MaybeUninit::write_slice_cloned(this, self.as_slice()); - new.init = self.init; - new - } -} - -impl Drop for Remainder { - fn drop(&mut self) { - // SAFETY: This raw slice will only contain the initialized objects - // within the buffer. - unsafe { ptr::drop_in_place(self.as_mut_slice()) } - } -} - /// 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 @@ -75,7 +17,7 @@ impl Drop for Remainder { #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] pub struct ArrayChunks { iter: Fuse, - remainder: Remainder, + remainder: Option>, } impl ArrayChunks @@ -84,25 +26,16 @@ where { pub(in crate::iter) fn new(iter: I) -> Self { assert!(N != 0, "chunk size must be non-zero"); - Self { iter: iter.fuse(), remainder: Remainder::new() } - } - - /// Returns a reference to the remaining elements of the original iterator - /// that are not going to be returned by this iterator. The returned slice - /// has at most `N-1` elements. - #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] - #[inline] - pub fn remainder(&self) -> &[I::Item] { - self.remainder.as_slice() + Self { iter: iter.fuse(), remainder: None } } - /// Returns a mutable reference to the remaining elements of the original - /// iterator that are not going to be returned by this iterator. The - /// returned slice has at most `N-1` elements. + /// 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 = "none")] #[inline] - pub fn remainder_mut(&mut self) -> &mut [I::Item] { - self.remainder.as_mut_slice() + pub fn into_remainder(self) -> Option> { + self.remainder } } @@ -129,8 +62,10 @@ where if guard.init > 0 { let init = guard.init; mem::forget(guard); - // SAFETY: `array` was initialized with `init` elements. - self.remainder = unsafe { Remainder::with_init(array, init) }; + self.remainder = { + // SAFETY: `array` was initialized with `init` elements. + Some(unsafe { array::IntoIter::with_partial(array, 0..init) }) + }; } return None; } @@ -189,7 +124,7 @@ where let init = guard.init; mem::forget(guard); // SAFETY: `array` was initialized with `init` elements. - self.remainder = unsafe { Remainder::with_init(array, init) }; + self.remainder = Some(unsafe { array::IntoIter::with_partial(array, 0..init) }); } R::from_output(o) } @@ -370,7 +305,7 @@ where // SAFETY: `array` was initialized with exactly `init` elements. self.remainder = unsafe { array.get_unchecked_mut(..init).reverse(); - Remainder::with_init(array, init) + Some(array::IntoIter::with_partial(array, 0..init)) }; Some(()) } diff --git a/core/src/iter/traits/iterator.rs b/core/src/iter/traits/iterator.rs index 8bf41ca6f..d41cf78e0 100644 --- a/core/src/iter/traits/iterator.rs +++ b/core/src/iter/traits/iterator.rs @@ -3319,7 +3319,9 @@ pub trait Iterator { /// 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. + /// 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 /// @@ -3336,7 +3338,7 @@ pub trait Iterator { /// assert_eq!(iter.next(), Some(['l', 'o'])); /// assert_eq!(iter.next(), Some(['r', 'e'])); /// assert_eq!(iter.next(), None); - /// assert_eq!(iter.remainder(), &['m']); + /// assert_eq!(iter.into_remainder().unwrap().as_slice(), &['m']); /// ``` /// /// ``` diff --git a/core/tests/iter/adapters/array_chunks.rs b/core/tests/iter/adapters/array_chunks.rs index 6845c94d3..dbcfd4560 100644 --- a/core/tests/iter/adapters/array_chunks.rs +++ b/core/tests/iter/adapters/array_chunks.rs @@ -15,39 +15,24 @@ fn test_iterator_array_chunks_infer() { 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); - assert_eq!(it.remainder().len(), 2); - let mut it2 = it.clone(); assert_eq!(count.get(), 3); - assert_eq!(it2.remainder().len(), 2); - - drop(it); + assert_eq!(it.into_remainder().unwrap().len(), 2); assert_eq!(count.get(), 5); - assert_eq!(it2.remainder().len(), 2); assert!(it2.next().is_none()); - - drop(it2); + 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.remainder(), &[]); - assert_eq!(it.remainder_mut(), &[]); assert_eq!(it.next(), Some([0, 1, 2, 3])); - assert_eq!(it.remainder(), &[]); - assert_eq!(it.remainder_mut(), &[]); assert_eq!(it.next(), Some([4, 5, 6, 7])); - assert_eq!(it.remainder(), &[]); - assert_eq!(it.remainder_mut(), &[]); - assert_eq!(it.next(), None); assert_eq!(it.next(), None); - assert_eq!(it.remainder(), &[8, 9, 10]); - assert_eq!(it.remainder_mut(), &[8, 9, 10]); + assert_eq!(it.into_remainder().unwrap().as_slice(), &[8, 9, 10]); } #[test] @@ -105,8 +90,7 @@ fn test_iterator_array_chunks_next_and_next_back() { assert_eq!(it.next(), None); assert_eq!(it.next_back(), None); assert_eq!(it.next(), None); - assert_eq!(it.remainder(), &[9, 10]); - assert_eq!(it.remainder_mut(), &[9, 10]); + assert_eq!(it.into_remainder().unwrap().as_slice(), &[9, 10]); } #[test] @@ -119,7 +103,7 @@ fn test_iterator_array_chunks_rev_remainder() { assert_eq!(it.next(), None); assert_eq!(it.next(), None); } - assert_eq!(it.remainder(), &[8, 9, 10]); + assert_eq!(it.into_remainder().unwrap().as_slice(), &[8, 9, 10]); } #[test] @@ -128,7 +112,6 @@ fn test_iterator_array_chunks_try_fold() { 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!(it.remainder().len(), 1); assert_eq!(count.get(), 9); drop(it); assert_eq!(count.get(), 10); @@ -137,7 +120,6 @@ fn test_iterator_array_chunks_try_fold() { 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!(it.remainder().len(), 0); assert_eq!(count.get(), 9); drop(it); assert_eq!(count.get(), 9); @@ -166,7 +148,6 @@ fn test_iterator_array_chunks_try_rfold() { 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!(it.remainder().len(), 1); assert_eq!(count.get(), 9); drop(it); assert_eq!(count.get(), 10); From a6239e3af3c5fdb6f4ae1ce18804edb472117954 Mon Sep 17 00:00:00 2001 From: Maybe Waffle Date: Mon, 1 Aug 2022 17:00:51 +0400 Subject: [PATCH 03/12] Remove `array::IntoIter::with_partial` -- an artifact of the past, once used to create an `IntoIter` from its parts --- core/src/array/iter.rs | 10 ---------- core/src/iter/adapters/array_chunks.rs | 7 ++++--- 2 files changed, 4 insertions(+), 13 deletions(-) diff --git a/core/src/array/iter.rs b/core/src/array/iter.rs index 459cd094c..f4885ed9f 100644 --- a/core/src/array/iter.rs +++ b/core/src/array/iter.rs @@ -84,16 +84,6 @@ impl IntoIter { IntoIterator::into_iter(array) } - /// Creates a new iterator from a partially initalized array. - /// - /// # Safety - /// - /// The caller must guarantee that all and only the `alive` elements of - /// `data` are initialized. - pub(crate) unsafe fn with_partial(data: [MaybeUninit; N], alive: Range) -> Self { - Self { data, alive } - } - /// Creates an iterator over the elements in a partially-initialized buffer. /// /// If you have a fully-initialized array, then use [`IntoIterator`]. diff --git a/core/src/iter/adapters/array_chunks.rs b/core/src/iter/adapters/array_chunks.rs index 2ec1284c3..8f3e1b58b 100644 --- a/core/src/iter/adapters/array_chunks.rs +++ b/core/src/iter/adapters/array_chunks.rs @@ -64,7 +64,7 @@ where mem::forget(guard); self.remainder = { // SAFETY: `array` was initialized with `init` elements. - Some(unsafe { array::IntoIter::with_partial(array, 0..init) }) + Some(unsafe { array::IntoIter::new_unchecked(array, 0..init) }) }; } return None; @@ -124,7 +124,8 @@ where let init = guard.init; mem::forget(guard); // SAFETY: `array` was initialized with `init` elements. - self.remainder = Some(unsafe { array::IntoIter::with_partial(array, 0..init) }); + self.remainder = + Some(unsafe { array::IntoIter::new_unchecked(array, 0..init) }); } R::from_output(o) } @@ -305,7 +306,7 @@ where // SAFETY: `array` was initialized with exactly `init` elements. self.remainder = unsafe { array.get_unchecked_mut(..init).reverse(); - Some(array::IntoIter::with_partial(array, 0..init)) + Some(array::IntoIter::new_unchecked(array, 0..init)) }; Some(()) } From 55624f614acea2fbf3f5f722bebc2049875222f4 Mon Sep 17 00:00:00 2001 From: Maybe Waffle Date: Mon, 1 Aug 2022 18:26:18 +0400 Subject: [PATCH 04/12] Forward `ArrayChunks::next{,_back}` to `try_{for_each,rfold}` (suggested in the review of the previous attempt to add `ArrayChunks`) --- core/src/iter/adapters/array_chunks.rs | 44 ++------------------------ 1 file changed, 2 insertions(+), 42 deletions(-) diff --git a/core/src/iter/adapters/array_chunks.rs b/core/src/iter/adapters/array_chunks.rs index 8f3e1b58b..e25a6f975 100644 --- a/core/src/iter/adapters/array_chunks.rs +++ b/core/src/iter/adapters/array_chunks.rs @@ -48,33 +48,7 @@ where #[inline] fn next(&mut self) -> Option { - let mut array = MaybeUninit::uninit_array(); - // SAFETY: `array` will still be valid if `guard` is dropped. - let mut guard = unsafe { FrontGuard::new(&mut array) }; - - for slot in array.iter_mut() { - match self.iter.next() { - Some(item) => { - slot.write(item); - guard.init += 1; - } - None => { - if guard.init > 0 { - let init = guard.init; - mem::forget(guard); - self.remainder = { - // SAFETY: `array` was initialized with `init` elements. - Some(unsafe { array::IntoIter::new_unchecked(array, 0..init) }) - }; - } - return None; - } - } - } - - mem::forget(guard); - // SAFETY: All elements of the array were populated in the loop above. - Some(unsafe { MaybeUninit::array_assume_init(array) }) + self.try_for_each(ControlFlow::Break).break_value() } #[inline] @@ -194,21 +168,7 @@ where { #[inline] fn next_back(&mut self) -> Option { - // We are iterating from the back we need to first handle the remainder. - self.next_back_remainder()?; - - let mut array = MaybeUninit::uninit_array(); - // SAFETY: `array` will still be valid if `guard` is dropped. - let mut guard = unsafe { BackGuard::new(&mut array) }; - - for slot in array.iter_mut().rev() { - slot.write(self.iter.next_back()?); - guard.uninit -= 1; - } - - mem::forget(guard); - // SAFETY: All elements of the array were populated in the loop above. - Some(unsafe { MaybeUninit::array_assume_init(array) }) + self.try_rfold((), |(), x| ControlFlow::Break(x)).break_value() } fn try_rfold(&mut self, init: B, mut f: F) -> R From 587dcc5bd86c1c8fd627fc3802b2c04e2828fc1b Mon Sep 17 00:00:00 2001 From: Maybe Waffle Date: Mon, 1 Aug 2022 18:30:55 +0400 Subject: [PATCH 05/12] Remove incorrect impl `TrustedLen` for `ArrayChunks` As explained in the review of the previous attempt to add `ArrayChunks`, adapters that shrink the length can't implement `TrustedLen`. --- core/src/iter/adapters/array_chunks.rs | 11 ++--------- core/tests/iter/adapters/array_chunks.rs | 2 +- 2 files changed, 3 insertions(+), 10 deletions(-) diff --git a/core/src/iter/adapters/array_chunks.rs b/core/src/iter/adapters/array_chunks.rs index e25a6f975..c2de5efed 100644 --- a/core/src/iter/adapters/array_chunks.rs +++ b/core/src/iter/adapters/array_chunks.rs @@ -1,5 +1,5 @@ use crate::array; -use crate::iter::{Fuse, FusedIterator, Iterator, TrustedLen}; +use crate::iter::{Fuse, FusedIterator, Iterator}; use crate::mem; use crate::mem::MaybeUninit; use crate::ops::{ControlFlow, Try}; @@ -54,11 +54,7 @@ where #[inline] fn size_hint(&self) -> (usize, Option) { let (lower, upper) = self.iter.size_hint(); - // Keep infinite iterator size hint lower bound as `usize::MAX`. This - // is required to implement `TrustedLen`. - if lower == usize::MAX { - return (lower, upper); - } + (lower / N, upper.map(|n| n / N)) } @@ -318,6 +314,3 @@ where self.iter.len() / N == 0 } } - -#[unstable(feature = "trusted_len", issue = "37572")] -unsafe impl TrustedLen for ArrayChunks where I: TrustedLen {} diff --git a/core/tests/iter/adapters/array_chunks.rs b/core/tests/iter/adapters/array_chunks.rs index dbcfd4560..4e9d89e1e 100644 --- a/core/tests/iter/adapters/array_chunks.rs +++ b/core/tests/iter/adapters/array_chunks.rs @@ -50,7 +50,7 @@ fn test_iterator_array_chunks_size_hint() { assert_eq!(it.size_hint(), (0, Some(0))); let it = (1..).array_chunks::<2>(); - assert_eq!(it.size_hint(), (usize::MAX, None)); + 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)); From b65f0a70e4f23f4427db6c4534fc5284ea08adbb Mon Sep 17 00:00:00 2001 From: Maybe Waffle Date: Mon, 1 Aug 2022 18:34:30 +0400 Subject: [PATCH 06/12] Use `#[track_caller]` to make panic in `Iterator::array_chunks` nicer --- core/src/iter/adapters/array_chunks.rs | 1 + core/src/iter/traits/iterator.rs | 1 + 2 files changed, 2 insertions(+) diff --git a/core/src/iter/adapters/array_chunks.rs b/core/src/iter/adapters/array_chunks.rs index c2de5efed..f8a52ecb6 100644 --- a/core/src/iter/adapters/array_chunks.rs +++ b/core/src/iter/adapters/array_chunks.rs @@ -24,6 +24,7 @@ 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: iter.fuse(), remainder: None } diff --git a/core/src/iter/traits/iterator.rs b/core/src/iter/traits/iterator.rs index d41cf78e0..95c7cf575 100644 --- a/core/src/iter/traits/iterator.rs +++ b/core/src/iter/traits/iterator.rs @@ -3350,6 +3350,7 @@ pub trait Iterator { /// assert_eq!(x + y + z, 4); /// } /// ``` + #[track_caller] #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] fn array_chunks(self) -> ArrayChunks where From 4b8580cc2aa5642312a3b4f142056592aad19802 Mon Sep 17 00:00:00 2001 From: Maybe Waffle Date: Mon, 1 Aug 2022 18:43:40 +0400 Subject: [PATCH 07/12] Remove `Fuse` from `ArrayChunks` implementation It doesn't seem to be used at all. --- core/src/iter/adapters/array_chunks.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/core/src/iter/adapters/array_chunks.rs b/core/src/iter/adapters/array_chunks.rs index f8a52ecb6..3e8f6281d 100644 --- a/core/src/iter/adapters/array_chunks.rs +++ b/core/src/iter/adapters/array_chunks.rs @@ -1,5 +1,5 @@ use crate::array; -use crate::iter::{Fuse, FusedIterator, Iterator}; +use crate::iter::{FusedIterator, Iterator}; use crate::mem; use crate::mem::MaybeUninit; use crate::ops::{ControlFlow, Try}; @@ -16,7 +16,7 @@ use crate::ptr; #[must_use = "iterators are lazy and do nothing unless consumed"] #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] pub struct ArrayChunks { - iter: Fuse, + iter: I, remainder: Option>, } @@ -27,7 +27,7 @@ where #[track_caller] pub(in crate::iter) fn new(iter: I) -> Self { assert!(N != 0, "chunk size must be non-zero"); - Self { iter: iter.fuse(), remainder: None } + Self { iter, remainder: None } } /// Returns an iterator over the remaining elements of the original iterator From 02ec67bb9a61fd2791cfcf9888e2984396ccc31f Mon Sep 17 00:00:00 2001 From: Maybe Waffle Date: Mon, 1 Aug 2022 18:48:47 +0400 Subject: [PATCH 08/12] Simplify `ArrayChunks::is_empty` --- core/src/iter/adapters/array_chunks.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/src/iter/adapters/array_chunks.rs b/core/src/iter/adapters/array_chunks.rs index 3e8f6281d..901f559c4 100644 --- a/core/src/iter/adapters/array_chunks.rs +++ b/core/src/iter/adapters/array_chunks.rs @@ -312,6 +312,6 @@ where #[inline] fn is_empty(&self) -> bool { - self.iter.len() / N == 0 + self.iter.len() < N } } From de097a3a6c1bde3805b441c5c748828b33fed6f5 Mon Sep 17 00:00:00 2001 From: Maybe Waffle Date: Mon, 1 Aug 2022 19:06:13 +0400 Subject: [PATCH 09/12] Simplify `ArrayChunks::{,r}fold` impls --- core/src/iter/adapters/array_chunks.rs | 50 +++----------------------- 1 file changed, 4 insertions(+), 46 deletions(-) diff --git a/core/src/iter/adapters/array_chunks.rs b/core/src/iter/adapters/array_chunks.rs index 901f559c4..b66e23c1e 100644 --- a/core/src/iter/adapters/array_chunks.rs +++ b/core/src/iter/adapters/array_chunks.rs @@ -2,7 +2,7 @@ use crate::array; use crate::iter::{FusedIterator, Iterator}; use crate::mem; use crate::mem::MaybeUninit; -use crate::ops::{ControlFlow, Try}; +use crate::ops::{ControlFlow, NeverShortCircuit, Try}; use crate::ptr; /// An iterator over `N` elements of the iterator at a time. @@ -104,30 +104,12 @@ where } } - fn fold(self, init: B, mut f: F) -> B + fn fold(mut self, init: B, mut f: F) -> B where Self: Sized, F: FnMut(B, Self::Item) -> B, { - let mut array = MaybeUninit::uninit_array(); - // SAFETY: `array` will still be valid if `guard` is dropped. - let mut guard = unsafe { FrontGuard::new(&mut array) }; - - self.iter.fold(init, |mut acc, item| { - // SAFETY: `init` starts at 0, increases by one each iteration and - // is reset to 0 once it reaches N. - unsafe { array.get_unchecked_mut(guard.init) }.write(item); - guard.init += 1; - if guard.init == N { - guard.init = 0; - let array = mem::replace(&mut array, MaybeUninit::uninit_array()); - // SAFETY: the condition above asserts that all elements are - // initialized. - let item = unsafe { MaybeUninit::array_assume_init(array) }; - acc = f(acc, item); - } - acc - }) + self.try_fold(init, |acc, x| NeverShortCircuit(f(acc, x))).0 } } @@ -205,31 +187,7 @@ where Self: Sized, F: FnMut(B, Self::Item) -> B, { - // We are iterating from the back we need to first handle the remainder. - if self.next_back_remainder().is_none() { - return init; - } - - let mut array = MaybeUninit::uninit_array(); - - // SAFETY: `array` will still be valid if `guard` is dropped. - let mut guard = unsafe { BackGuard::new(&mut array) }; - - self.iter.rfold(init, |mut acc, item| { - guard.uninit -= 1; - // SAFETY: `uninit` starts at N, decreases by one each iteration and - // is reset to N once it reaches 0. - unsafe { array.get_unchecked_mut(guard.uninit) }.write(item); - if guard.uninit == 0 { - guard.uninit = N; - let array = mem::replace(&mut array, MaybeUninit::uninit_array()); - // SAFETY: the condition above asserts that all elements are - // initialized. - let item = unsafe { MaybeUninit::array_assume_init(array) }; - acc = f(acc, item); - } - acc - }) + self.try_rfold(init, |acc, x| NeverShortCircuit(f(acc, x))).0 } } From 69fc19af73b287815dfeed1aa16d2c9b2562a272 Mon Sep 17 00:00:00 2001 From: Maybe Waffle Date: Tue, 2 Aug 2022 10:42:16 +0400 Subject: [PATCH 10/12] Use `next_chunk` in `ArrayChunks` impl --- core/src/iter/adapters/array_chunks.rs | 169 ++++++------------------- 1 file changed, 37 insertions(+), 132 deletions(-) diff --git a/core/src/iter/adapters/array_chunks.rs b/core/src/iter/adapters/array_chunks.rs index b66e23c1e..3af72c16a 100644 --- a/core/src/iter/adapters/array_chunks.rs +++ b/core/src/iter/adapters/array_chunks.rs @@ -1,9 +1,6 @@ use crate::array; use crate::iter::{FusedIterator, Iterator}; -use crate::mem; -use crate::mem::MaybeUninit; use crate::ops::{ControlFlow, NeverShortCircuit, Try}; -use crate::ptr; /// An iterator over `N` elements of the iterator at a time. /// @@ -70,37 +67,18 @@ where F: FnMut(B, Self::Item) -> R, R: Try, { - let mut array = MaybeUninit::uninit_array(); - // SAFETY: `array` will still be valid if `guard` is dropped. - let mut guard = unsafe { FrontGuard::new(&mut array) }; - - let result = self.iter.try_fold(init, |mut acc, item| { - // SAFETY: `init` starts at 0, increases by one each iteration and - // is reset to 0 once it reaches N. - unsafe { array.get_unchecked_mut(guard.init) }.write(item); - guard.init += 1; - if guard.init == N { - guard.init = 0; - let array = mem::replace(&mut array, MaybeUninit::uninit_array()); - // SAFETY: the condition above asserts that all elements are - // initialized. - let item = unsafe { MaybeUninit::array_assume_init(array) }; - acc = f(acc, item)?; - } - R::from_output(acc) - }); - match result.branch() { - ControlFlow::Continue(o) => { - if guard.init > 0 { - let init = guard.init; - mem::forget(guard); - // SAFETY: `array` was initialized with `init` elements. - self.remainder = - Some(unsafe { array::IntoIter::new_unchecked(array, 0..init) }); + 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 }; } - R::from_output(o) } - ControlFlow::Break(r) => R::from_residual(r), } } @@ -113,33 +91,6 @@ where } } -/// A guard for an array where elements are filled from the left. -struct FrontGuard { - /// A pointer to the array that is being filled. We need to use a raw - /// pointer here because of the lifetime issues in the fold implementations. - ptr: *mut T, - /// The number of *initialized* elements. - init: usize, -} - -impl FrontGuard { - unsafe fn new(array: &mut [MaybeUninit; N]) -> Self { - Self { ptr: MaybeUninit::slice_as_mut_ptr(array), init: 0 } - } -} - -impl Drop for FrontGuard { - fn drop(&mut self) { - debug_assert!(self.init <= N); - // SAFETY: This raw slice will only contain the initialized objects - // within the buffer. - unsafe { - let slice = ptr::slice_from_raw_parts_mut(self.ptr, self.init); - ptr::drop_in_place(slice); - } - } -} - #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] impl DoubleEndedIterator for ArrayChunks where @@ -157,29 +108,20 @@ where R: Try, { // We are iterating from the back we need to first handle the remainder. - if self.next_back_remainder().is_none() { - return R::from_output(init); - } + self.next_back_remainder(); - let mut array = MaybeUninit::uninit_array(); - // SAFETY: `array` will still be valid if `guard` is dropped. - let mut guard = unsafe { BackGuard::new(&mut array) }; + let mut acc = init; + let mut iter = self.iter.by_ref().rev(); - self.iter.try_rfold(init, |mut acc, item| { - guard.uninit -= 1; - // SAFETY: `uninit` starts at N, decreases by one each iteration and - // is reset to N once it reaches 0. - unsafe { array.get_unchecked_mut(guard.uninit) }.write(item); - if guard.uninit == 0 { - guard.uninit = N; - let array = mem::replace(&mut array, MaybeUninit::uninit_array()); - // SAFETY: the condition above asserts that all elements are - // initialized. - let item = unsafe { MaybeUninit::array_assume_init(array) }; - acc = f(acc, item)?; - } - R::from_output(acc) - }) + // 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() { + chunk.reverse(); + acc = f(acc, chunk)? + } + + try { acc } } fn rfold(mut self, init: B, mut f: F) -> B @@ -195,63 +137,26 @@ impl ArrayChunks where I: DoubleEndedIterator + ExactSizeIterator, { - #[inline] - fn next_back_remainder(&mut self) -> Option<()> { + /// 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; - if rem == 0 { - return Some(()); - } - - let mut array = MaybeUninit::uninit_array(); - // SAFETY: The array will still be valid if `guard` is dropped and - // it is forgotten otherwise. - let mut guard = unsafe { FrontGuard::new(&mut array) }; + // 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() }; - // SAFETY: `rem` is in the range 1..N based on how it is calculated. - for slot in unsafe { array.get_unchecked_mut(..rem) }.iter_mut() { - slot.write(self.iter.next_back()?); - guard.init += 1; - } - - let init = guard.init; - mem::forget(guard); - // SAFETY: `array` was initialized with exactly `init` elements. - self.remainder = unsafe { - array.get_unchecked_mut(..init).reverse(); - Some(array::IntoIter::new_unchecked(array, 0..init)) - }; - Some(()) - } -} - -/// A guard for an array where elements are filled from the right. -struct BackGuard { - /// A pointer to the array that is being filled. We need to use a raw - /// pointer here because of the lifetime issues in the rfold implementations. - ptr: *mut T, - /// The number of *uninitialized* elements. - uninit: usize, -} - -impl BackGuard { - unsafe fn new(array: &mut [MaybeUninit; N]) -> Self { - Self { ptr: MaybeUninit::slice_as_mut_ptr(array), uninit: N } - } -} - -impl Drop for BackGuard { - fn drop(&mut self) { - debug_assert!(self.uninit <= N); - // SAFETY: This raw slice will only contain the initialized objects - // within the buffer. - unsafe { - let ptr = self.ptr.offset(self.uninit as isize); - let slice = ptr::slice_from_raw_parts_mut(ptr, N - self.uninit); - ptr::drop_in_place(slice); - } + // We used `.rev()` above, so we need to re-reverse the reminder + remainder.as_mut_slice().reverse(); + self.remainder = Some(remainder); } } From 7611d2209b6517f6f9d3cdc0b007028129746305 Mon Sep 17 00:00:00 2001 From: Maybe Waffle Date: Fri, 12 Aug 2022 14:57:15 +0400 Subject: [PATCH 11/12] address review comments --- core/src/iter/adapters/array_chunks.rs | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/core/src/iter/adapters/array_chunks.rs b/core/src/iter/adapters/array_chunks.rs index 3af72c16a..86835f443 100644 --- a/core/src/iter/adapters/array_chunks.rs +++ b/core/src/iter/adapters/array_chunks.rs @@ -1,5 +1,5 @@ use crate::array; -use crate::iter::{FusedIterator, Iterator}; +use crate::iter::{ByRefSized, FusedIterator, Iterator}; use crate::ops::{ControlFlow, NeverShortCircuit, Try}; /// An iterator over `N` elements of the iterator at a time. @@ -82,12 +82,12 @@ where } } - fn fold(mut self, init: B, mut f: F) -> B + fn fold(mut self, init: B, f: F) -> B where Self: Sized, F: FnMut(B, Self::Item) -> B, { - self.try_fold(init, |acc, x| NeverShortCircuit(f(acc, x))).0 + self.try_fold(init, NeverShortCircuit::wrap_mut_2(f)).0 } } @@ -111,12 +111,14 @@ where self.next_back_remainder(); let mut acc = init; - let mut iter = self.iter.by_ref().rev(); + 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)? } @@ -124,12 +126,12 @@ where try { acc } } - fn rfold(mut self, init: B, mut f: F) -> B + fn rfold(mut self, init: B, f: F) -> B where Self: Sized, F: FnMut(B, Self::Item) -> B, { - self.try_rfold(init, |acc, x| NeverShortCircuit(f(acc, x))).0 + self.try_rfold(init, NeverShortCircuit::wrap_mut_2(f)).0 } } From 1b5d97f87f9ec189bc3719516aa27df70aaae8a1 Mon Sep 17 00:00:00 2001 From: Maybe Waffle Date: Fri, 12 Aug 2022 14:58:14 +0400 Subject: [PATCH 12/12] fill-in tracking issue for `feature(iter_array_chunks)` --- core/src/iter/adapters/array_chunks.rs | 12 ++++++------ core/src/iter/adapters/mod.rs | 2 +- core/src/iter/mod.rs | 2 +- core/src/iter/traits/iterator.rs | 2 +- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/core/src/iter/adapters/array_chunks.rs b/core/src/iter/adapters/array_chunks.rs index 86835f443..9b479a9f8 100644 --- a/core/src/iter/adapters/array_chunks.rs +++ b/core/src/iter/adapters/array_chunks.rs @@ -11,7 +11,7 @@ use crate::ops::{ControlFlow, NeverShortCircuit, Try}; /// 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 = "none")] +#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")] pub struct ArrayChunks { iter: I, remainder: Option>, @@ -30,14 +30,14 @@ where /// 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 = "none")] + #[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 = "none")] +#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")] impl Iterator for ArrayChunks where I: Iterator, @@ -91,7 +91,7 @@ where } } -#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] +#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")] impl DoubleEndedIterator for ArrayChunks where I: DoubleEndedIterator + ExactSizeIterator, @@ -162,10 +162,10 @@ where } } -#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] +#[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 = "none")] +#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")] impl ExactSizeIterator for ArrayChunks where I: ExactSizeIterator, diff --git a/core/src/iter/adapters/mod.rs b/core/src/iter/adapters/mod.rs index 39e7ab878..bf4fabad3 100644 --- a/core/src/iter/adapters/mod.rs +++ b/core/src/iter/adapters/mod.rs @@ -33,7 +33,7 @@ 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 = "none")] +#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")] pub use self::array_chunks::ArrayChunks; #[unstable(feature = "std_internals", issue = "none")] diff --git a/core/src/iter/mod.rs b/core/src/iter/mod.rs index d48e3a52c..9514466bd 100644 --- a/core/src/iter/mod.rs +++ b/core/src/iter/mod.rs @@ -398,7 +398,7 @@ 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 = "none")] +#[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; diff --git a/core/src/iter/traits/iterator.rs b/core/src/iter/traits/iterator.rs index 95c7cf575..b2d08f4b0 100644 --- a/core/src/iter/traits/iterator.rs +++ b/core/src/iter/traits/iterator.rs @@ -3351,7 +3351,7 @@ pub trait Iterator { /// } /// ``` #[track_caller] - #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] + #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")] fn array_chunks(self) -> ArrayChunks where Self: Sized,