From 563db25776c51af11595628ee706c47c8ea94735 Mon Sep 17 00:00:00 2001 From: John Schug Date: Wed, 4 Aug 2021 04:21:27 -0700 Subject: [PATCH 1/8] Add Iterator::{array_chunks, array_rchunks} methods --- .../core/src/iter/adapters/array_chunks.rs | 146 ++++++++++++++++++ library/core/src/iter/adapters/mod.rs | 4 + library/core/src/iter/mod.rs | 2 + library/core/src/iter/traits/iterator.rs | 61 ++++++++ .../core/tests/iter/adapters/array_chunks.rs | 114 ++++++++++++++ library/core/tests/iter/adapters/mod.rs | 1 + library/core/tests/lib.rs | 1 + 7 files changed, 329 insertions(+) create mode 100644 library/core/src/iter/adapters/array_chunks.rs create mode 100644 library/core/tests/iter/adapters/array_chunks.rs 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..4ab9c690f7ddf --- /dev/null +++ b/library/core/src/iter/adapters/array_chunks.rs @@ -0,0 +1,146 @@ +use core::iter::FusedIterator; +use core::mem::{self, MaybeUninit}; +use core::ptr; + +/// 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)] +pub struct ArrayChunks { + iter: I, + buffer: [MaybeUninit; N], + init: usize, +} + +#[unstable(feature = "iter_array_chunks", issue = "none")] +impl Drop for ArrayChunks { + fn drop(&mut self) { + // SAFETY: This is safe: `remainder_mut` returns exactly the sub-slice + // of elements that were initialized but not yielded and so have yet + // to be dropped. + unsafe { + ptr::drop_in_place(self.remainder_mut()); + } + } +} + +impl ArrayChunks { + pub(in crate::iter) fn new(iter: I) -> Self { + Self { iter, init: 0, buffer: MaybeUninit::uninit_array() } + } + + /// 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[..self.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[..self.init]) } + } +} + +#[unstable(feature = "iter_array_chunks", issue = "none")] +impl Iterator for ArrayChunks { + type Item = [I::Item; N]; + + fn next(&mut self) -> Option { + while self.init < N { + self.buffer[self.init] = MaybeUninit::new(self.iter.next()?); + self.init += 1; + } + self.init = 0; + // SAFETY: This is safe: `MaybeUninit` is guaranteed to have the same layout + // as `T` and the entire array has just been initialized. + unsafe { Some(mem::transmute_copy(&self.buffer)) } + } +} + +#[unstable(feature = "iter_array_chunks", issue = "none")] +impl FusedIterator for ArrayChunks {} + +/// An iterator that yields the elements of another iterator in +/// chunks of size `N` starting from the end. +/// +/// This `struct` is created by the [`array_rchunks`] method on [`Iterator`]. See +/// its documentation for more. +/// +/// [`array_rchunks`]: Iterator::array_rchunks +#[unstable(feature = "iter_array_chunks", issue = "none")] +#[derive(Debug)] +pub struct ArrayRChunks { + iter: I, + buffer: [MaybeUninit; N], + init: usize, +} + +#[unstable(feature = "iter_array_chunks", issue = "none")] +impl Drop for ArrayRChunks { + fn drop(&mut self) { + // SAFETY: This is safe: `remainder_mut` returns exactly the sub-slice + // of elements that were initialized but not yielded and so have yet + // to be dropped. + unsafe { + ptr::drop_in_place(self.remainder_mut()); + } + } +} + +impl ArrayRChunks { + pub(in crate::iter) fn new(iter: I) -> Self { + Self { iter, init: 0, buffer: MaybeUninit::uninit_array() } + } + + /// 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 after `init` are properly initialized. + unsafe { MaybeUninit::slice_assume_init_ref(&self.buffer[(N - self.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 after `init` are properly initialized. + unsafe { MaybeUninit::slice_assume_init_mut(&mut self.buffer[(N - self.init)..]) } + } +} + +#[unstable(feature = "iter_array_chunks", issue = "none")] +impl Iterator for ArrayRChunks { + type Item = [I::Item; N]; + + fn next(&mut self) -> Option { + while self.init < N { + self.buffer[N - self.init - 1] = MaybeUninit::new(self.iter.next_back()?); + self.init += 1; + } + self.init = 0; + // SAFETY: This is safe: `MaybeUninit` is guaranteed to have the same layout + // as `T` and the entire array has just been initialized. + unsafe { Some(mem::transmute_copy(&self.buffer)) } + } +} + +#[unstable(feature = "iter_array_chunks", issue = "none")] +impl FusedIterator for ArrayRChunks where + I: DoubleEndedIterator + FusedIterator +{ +} diff --git a/library/core/src/iter/adapters/mod.rs b/library/core/src/iter/adapters/mod.rs index 056ccca1d01c9..dc77d0a755d40 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::{ControlFlow, Try}; +mod array_chunks; mod chain; mod cloned; mod copied; @@ -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, ArrayRChunks}; + #[stable(feature = "iter_cloned", since = "1.1.0")] pub use self::cloned::Cloned; diff --git a/library/core/src/iter/mod.rs b/library/core/src/iter/mod.rs index 7fb80f954ff40..2567857b84e5f 100644 --- a/library/core/src/iter/mod.rs +++ b/library/core/src/iter/mod.rs @@ -409,6 +409,8 @@ pub use self::adapters::StepBy; pub use self::adapters::TrustedRandomAccess; #[unstable(feature = "trusted_random_access", issue = "none")] pub use self::adapters::TrustedRandomAccessNoCoerce; +#[unstable(feature = "iter_array_chunks", issue = "none")] +pub use self::adapters::{ArrayChunks, ArrayRChunks}; #[stable(feature = "rust1", since = "1.0.0")] pub use self::adapters::{ Chain, Cycle, Enumerate, Filter, FilterMap, FlatMap, Fuse, Inspect, Map, Peekable, Rev, Scan, diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs index 537e42f66de1b..fe403a1edfebb 100644 --- a/library/core/src/iter/traits/iterator.rs +++ b/library/core/src/iter/traits/iterator.rs @@ -6,6 +6,7 @@ use crate::cmp::{self, Ordering}; use crate::ops::{ControlFlow, Try}; use super::super::TrustedRandomAccessNoCoerce; +use super::super::{ArrayChunks, ArrayRChunks}; use super::super::{Chain, Cloned, Copied, Cycle, Enumerate, Filter, FilterMap, Fuse}; use super::super::{FlatMap, Flatten}; use super::super::{FromIterator, Intersperse, IntersperseWith, Product, Sum, Zip}; @@ -3468,6 +3469,66 @@ 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(self) -> ArrayChunks + where + Self: Sized, + { + assert_ne!(N, 0); + ArrayChunks::new(self) + } + + /// Creates an iterator which yields arrays of `N` elements yielded by + /// the original iterator starting from the end. + /// + /// # Panics + /// + /// Panics if `N` is zero. + /// + /// # Examples + /// + /// ``` + /// #![feature(iter_array_chunks)] + /// let mut iter = (0..10).array_rchunks::<3>(); + /// + /// assert_eq!(iter.next(), Some([7, 8, 9])); + /// assert_eq!(iter.next(), Some([4, 5, 6])); + /// assert_eq!(iter.next(), Some([1, 2, 3])); + /// assert_eq!(iter.next(), None); + /// + /// assert_eq!(iter.remainder(), &[0]); + /// ``` + #[inline] + #[unstable(feature = "iter_array_chunks", issue = "none")] + fn array_rchunks(self) -> ArrayRChunks + where + Self: Sized + DoubleEndedIterator, + { + assert_ne!(N, 0); + ArrayRChunks::new(self) + } } #[stable(feature = "rust1", since = "1.0.0")] 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..657de5bb99ac8 --- /dev/null +++ b/library/core/tests/iter/adapters/array_chunks.rs @@ -0,0 +1,114 @@ +use core::cell::Cell; +use core::iter::*; + +#[derive(Debug)] +struct DropBomb<'a> { + dropped: bool, + counter: &'a Cell, +} + +impl Drop for DropBomb<'_> { + fn drop(&mut self) { + if self.dropped { + panic!("double dropped!"); + } + self.dropped = true; + self.counter.set(self.counter.get() + 1); + } +} + +#[test] +fn test_iterator_array_chunks_remainder() { + let mut iter = (0..=10).array_chunks::<4>(); + assert_eq!(iter.remainder(), &[]); + assert_eq!(iter.remainder_mut(), &[]); + assert_eq!(iter.next(), Some([0, 1, 2, 3])); + assert_eq!(iter.remainder(), &[]); + assert_eq!(iter.remainder_mut(), &[]); + assert_eq!(iter.next(), Some([4, 5, 6, 7])); + assert_eq!(iter.remainder(), &[]); + assert_eq!(iter.remainder_mut(), &[]); + assert_eq!(iter.next(), None); + assert_eq!(iter.remainder(), &[8, 9, 10]); + assert_eq!(iter.remainder_mut(), &[8, 9, 10]); +} + +#[test] +fn test_iterator_array_chunks_drop() { + let counter = Cell::new(0); + let create = + |n| (0..n).map(|_| DropBomb { dropped: false, counter: &counter }).array_chunks::<3>(); + + let iter = create(5); + assert_eq!(counter.get(), 0); + drop(iter); + assert_eq!(counter.get(), 0); + + let mut iter = create(3); + counter.set(0); + iter.next(); + assert_eq!(counter.get(), 3); + assert!(iter.next().is_none()); + assert_eq!(counter.get(), 3); + assert_eq!(iter.remainder().len(), 0); + drop(iter); + assert_eq!(counter.get(), 3); + + let mut iter = create(5); + counter.set(0); + iter.next(); + assert_eq!(counter.get(), 3); + assert!(iter.next().is_none()); + assert_eq!(counter.get(), 3); + assert_eq!(iter.remainder().len(), 2); + drop(iter); + assert_eq!(counter.get(), 5); +} + +#[test] +fn test_iterator_array_rchunks_remainder() { + let mut iter = (0..=10).array_rchunks::<4>(); + assert_eq!(iter.remainder(), &[]); + assert_eq!(iter.remainder_mut(), &[]); + assert_eq!(iter.next(), Some([7, 8, 9, 10])); + assert_eq!(iter.remainder(), &[]); + assert_eq!(iter.remainder_mut(), &[]); + assert_eq!(iter.next(), Some([3, 4, 5, 6])); + assert_eq!(iter.remainder(), &[]); + assert_eq!(iter.remainder_mut(), &[]); + assert_eq!(iter.next(), None); + assert_eq!(iter.remainder(), &[0, 1, 2]); + assert_eq!(iter.remainder_mut(), &[0, 1, 2]); +} + +#[test] +fn test_iterator_array_rchunks_drop() { + let counter = Cell::new(0); + let create = + |n| (0..n).map(|_| DropBomb { dropped: false, counter: &counter }).array_rchunks::<3>(); + + let iter = create(5); + assert_eq!(counter.get(), 0); + drop(iter); + assert_eq!(counter.get(), 0); + + let mut iter = create(3); + counter.set(0); + iter.next(); + assert_eq!(counter.get(), 3); + assert!(iter.next().is_none()); + assert_eq!(counter.get(), 3); + assert_eq!(iter.remainder().len(), 0); + drop(iter); + assert_eq!(counter.get(), 3); + + let mut iter = create(5); + counter.set(0); + iter.next(); + assert_eq!(counter.get(), 3); + assert!(iter.next().is_none()); + assert_eq!(counter.get(), 3); + assert_eq!(iter.remainder().len(), 2); + drop(iter); + assert_eq!(counter.get(), 5); +} diff --git a/library/core/tests/iter/adapters/mod.rs b/library/core/tests/iter/adapters/mod.rs index 567d9fe49cade..a546e72a7b737 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; diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index c7756a503c3e9..380ac8ec26af5 100644 --- a/library/core/tests/lib.rs +++ b/library/core/tests/lib.rs @@ -45,6 +45,7 @@ #![feature(slice_partition_dedup)] #![feature(int_log)] #![feature(iter_advance_by)] +#![feature(iter_array_chunks)] #![feature(iter_partition_in_place)] #![feature(iter_intersperse)] #![feature(iter_is_partitioned)] From 4f507b98b98550e7de81ea9af97742c13074762f Mon Sep 17 00:00:00 2001 From: John Schug Date: Thu, 5 Aug 2021 16:34:24 -0700 Subject: [PATCH 2/8] Remove somewhat redundant array_rchunks method --- .../core/src/iter/adapters/array_chunks.rs | 73 ------------------- library/core/src/iter/adapters/mod.rs | 2 +- library/core/src/iter/mod.rs | 4 +- library/core/src/iter/traits/iterator.rs | 32 +------- .../core/tests/iter/adapters/array_chunks.rs | 48 ------------ 5 files changed, 4 insertions(+), 155 deletions(-) diff --git a/library/core/src/iter/adapters/array_chunks.rs b/library/core/src/iter/adapters/array_chunks.rs index 4ab9c690f7ddf..1043f200543e7 100644 --- a/library/core/src/iter/adapters/array_chunks.rs +++ b/library/core/src/iter/adapters/array_chunks.rs @@ -71,76 +71,3 @@ impl Iterator for ArrayChunks { #[unstable(feature = "iter_array_chunks", issue = "none")] impl FusedIterator for ArrayChunks {} - -/// An iterator that yields the elements of another iterator in -/// chunks of size `N` starting from the end. -/// -/// This `struct` is created by the [`array_rchunks`] method on [`Iterator`]. See -/// its documentation for more. -/// -/// [`array_rchunks`]: Iterator::array_rchunks -#[unstable(feature = "iter_array_chunks", issue = "none")] -#[derive(Debug)] -pub struct ArrayRChunks { - iter: I, - buffer: [MaybeUninit; N], - init: usize, -} - -#[unstable(feature = "iter_array_chunks", issue = "none")] -impl Drop for ArrayRChunks { - fn drop(&mut self) { - // SAFETY: This is safe: `remainder_mut` returns exactly the sub-slice - // of elements that were initialized but not yielded and so have yet - // to be dropped. - unsafe { - ptr::drop_in_place(self.remainder_mut()); - } - } -} - -impl ArrayRChunks { - pub(in crate::iter) fn new(iter: I) -> Self { - Self { iter, init: 0, buffer: MaybeUninit::uninit_array() } - } - - /// 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 after `init` are properly initialized. - unsafe { MaybeUninit::slice_assume_init_ref(&self.buffer[(N - self.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 after `init` are properly initialized. - unsafe { MaybeUninit::slice_assume_init_mut(&mut self.buffer[(N - self.init)..]) } - } -} - -#[unstable(feature = "iter_array_chunks", issue = "none")] -impl Iterator for ArrayRChunks { - type Item = [I::Item; N]; - - fn next(&mut self) -> Option { - while self.init < N { - self.buffer[N - self.init - 1] = MaybeUninit::new(self.iter.next_back()?); - self.init += 1; - } - self.init = 0; - // SAFETY: This is safe: `MaybeUninit` is guaranteed to have the same layout - // as `T` and the entire array has just been initialized. - unsafe { Some(mem::transmute_copy(&self.buffer)) } - } -} - -#[unstable(feature = "iter_array_chunks", issue = "none")] -impl FusedIterator for ArrayRChunks where - I: DoubleEndedIterator + FusedIterator -{ -} diff --git a/library/core/src/iter/adapters/mod.rs b/library/core/src/iter/adapters/mod.rs index dc77d0a755d40..d403a2d7a4c2b 100644 --- a/library/core/src/iter/adapters/mod.rs +++ b/library/core/src/iter/adapters/mod.rs @@ -32,7 +32,7 @@ pub use self::{ }; #[unstable(feature = "iter_array_chunks", issue = "none")] -pub use self::array_chunks::{ArrayChunks, ArrayRChunks}; +pub use self::array_chunks::ArrayChunks; #[stable(feature = "iter_cloned", since = "1.1.0")] pub use self::cloned::Cloned; diff --git a/library/core/src/iter/mod.rs b/library/core/src/iter/mod.rs index 2567857b84e5f..7232c9fcf303b 100644 --- a/library/core/src/iter/mod.rs +++ b/library/core/src/iter/mod.rs @@ -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")] @@ -409,8 +411,6 @@ pub use self::adapters::StepBy; pub use self::adapters::TrustedRandomAccess; #[unstable(feature = "trusted_random_access", issue = "none")] pub use self::adapters::TrustedRandomAccessNoCoerce; -#[unstable(feature = "iter_array_chunks", issue = "none")] -pub use self::adapters::{ArrayChunks, ArrayRChunks}; #[stable(feature = "rust1", since = "1.0.0")] pub use self::adapters::{ Chain, Cycle, Enumerate, Filter, FilterMap, FlatMap, Fuse, Inspect, Map, Peekable, Rev, Scan, diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs index fe403a1edfebb..961af58957822 100644 --- a/library/core/src/iter/traits/iterator.rs +++ b/library/core/src/iter/traits/iterator.rs @@ -5,8 +5,8 @@ use crate::cmp::{self, Ordering}; use crate::ops::{ControlFlow, Try}; +use super::super::ArrayChunks; use super::super::TrustedRandomAccessNoCoerce; -use super::super::{ArrayChunks, ArrayRChunks}; use super::super::{Chain, Cloned, Copied, Cycle, Enumerate, Filter, FilterMap, Fuse}; use super::super::{FlatMap, Flatten}; use super::super::{FromIterator, Intersperse, IntersperseWith, Product, Sum, Zip}; @@ -3499,36 +3499,6 @@ pub trait Iterator { assert_ne!(N, 0); ArrayChunks::new(self) } - - /// Creates an iterator which yields arrays of `N` elements yielded by - /// the original iterator starting from the end. - /// - /// # Panics - /// - /// Panics if `N` is zero. - /// - /// # Examples - /// - /// ``` - /// #![feature(iter_array_chunks)] - /// let mut iter = (0..10).array_rchunks::<3>(); - /// - /// assert_eq!(iter.next(), Some([7, 8, 9])); - /// assert_eq!(iter.next(), Some([4, 5, 6])); - /// assert_eq!(iter.next(), Some([1, 2, 3])); - /// assert_eq!(iter.next(), None); - /// - /// assert_eq!(iter.remainder(), &[0]); - /// ``` - #[inline] - #[unstable(feature = "iter_array_chunks", issue = "none")] - fn array_rchunks(self) -> ArrayRChunks - where - Self: Sized + DoubleEndedIterator, - { - assert_ne!(N, 0); - ArrayRChunks::new(self) - } } #[stable(feature = "rust1", since = "1.0.0")] diff --git a/library/core/tests/iter/adapters/array_chunks.rs b/library/core/tests/iter/adapters/array_chunks.rs index 657de5bb99ac8..67cb7ff44f5ff 100644 --- a/library/core/tests/iter/adapters/array_chunks.rs +++ b/library/core/tests/iter/adapters/array_chunks.rs @@ -64,51 +64,3 @@ fn test_iterator_array_chunks_drop() { drop(iter); assert_eq!(counter.get(), 5); } - -#[test] -fn test_iterator_array_rchunks_remainder() { - let mut iter = (0..=10).array_rchunks::<4>(); - assert_eq!(iter.remainder(), &[]); - assert_eq!(iter.remainder_mut(), &[]); - assert_eq!(iter.next(), Some([7, 8, 9, 10])); - assert_eq!(iter.remainder(), &[]); - assert_eq!(iter.remainder_mut(), &[]); - assert_eq!(iter.next(), Some([3, 4, 5, 6])); - assert_eq!(iter.remainder(), &[]); - assert_eq!(iter.remainder_mut(), &[]); - assert_eq!(iter.next(), None); - assert_eq!(iter.remainder(), &[0, 1, 2]); - assert_eq!(iter.remainder_mut(), &[0, 1, 2]); -} - -#[test] -fn test_iterator_array_rchunks_drop() { - let counter = Cell::new(0); - let create = - |n| (0..n).map(|_| DropBomb { dropped: false, counter: &counter }).array_rchunks::<3>(); - - let iter = create(5); - assert_eq!(counter.get(), 0); - drop(iter); - assert_eq!(counter.get(), 0); - - let mut iter = create(3); - counter.set(0); - iter.next(); - assert_eq!(counter.get(), 3); - assert!(iter.next().is_none()); - assert_eq!(counter.get(), 3); - assert_eq!(iter.remainder().len(), 0); - drop(iter); - assert_eq!(counter.get(), 3); - - let mut iter = create(5); - counter.set(0); - iter.next(); - assert_eq!(counter.get(), 3); - assert!(iter.next().is_none()); - assert_eq!(counter.get(), 3); - assert_eq!(iter.remainder().len(), 2); - drop(iter); - assert_eq!(counter.get(), 5); -} From 45bfc7fe2be4686e996136c6d55610ab1171b0ec Mon Sep 17 00:00:00 2001 From: John Schug Date: Thu, 5 Aug 2021 16:55:52 -0700 Subject: [PATCH 3/8] Override more Iterator methods in ArrayChunks --- .../core/src/iter/adapters/array_chunks.rs | 142 +++++++++++++++--- .../core/tests/iter/adapters/array_chunks.rs | 100 +++++++++++- 2 files changed, 214 insertions(+), 28 deletions(-) diff --git a/library/core/src/iter/adapters/array_chunks.rs b/library/core/src/iter/adapters/array_chunks.rs index 1043f200543e7..4e726579ad9b0 100644 --- a/library/core/src/iter/adapters/array_chunks.rs +++ b/library/core/src/iter/adapters/array_chunks.rs @@ -1,7 +1,54 @@ use core::iter::FusedIterator; use core::mem::{self, MaybeUninit}; +use core::ops::{ControlFlow, Try}; use core::ptr; +#[derive(Debug)] +struct Guard { + arr: [MaybeUninit; N], + init: usize, +} + +impl Guard { + fn new() -> Self { + Self { arr: MaybeUninit::uninit_array(), init: 0 } + } + + // SAFETY: The entire array must have been properly initialized. + unsafe fn assume_init_read(&mut self) -> [T; N] { + debug_assert_eq!(self.init, N); + // Reset the tracked initialization state. + self.init = 0; + // SAFETY: This is safe: `MaybeUninit` is guaranteed to have the same layout + // as `T` and per the function's contract the entire array has been initialized. + unsafe { mem::transmute_copy(&self.arr) } + } +} + +impl Clone for Guard { + 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.arr[..self.init]) }; + MaybeUninit::write_slice_cloned(&mut new.arr[..self.init], src); + new.init = self.init; + new + } +} + +impl Drop for Guard { + fn drop(&mut self) { + debug_assert!(self.init <= N); + + let init_part = &mut self.arr[..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(init_part)); + } + } +} + /// An iterator that yields the elements of another iterator in /// chunks of size `N`. /// @@ -10,28 +57,15 @@ use core::ptr; /// /// [`array_chunks`]: Iterator::array_chunks #[unstable(feature = "iter_array_chunks", issue = "none")] -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct ArrayChunks { iter: I, - buffer: [MaybeUninit; N], - init: usize, -} - -#[unstable(feature = "iter_array_chunks", issue = "none")] -impl Drop for ArrayChunks { - fn drop(&mut self) { - // SAFETY: This is safe: `remainder_mut` returns exactly the sub-slice - // of elements that were initialized but not yielded and so have yet - // to be dropped. - unsafe { - ptr::drop_in_place(self.remainder_mut()); - } - } + guard: Guard, } impl ArrayChunks { pub(in crate::iter) fn new(iter: I) -> Self { - Self { iter, init: 0, buffer: MaybeUninit::uninit_array() } + Self { iter, guard: Guard::new() } } /// Returns the remainder of the elements yielded by the original @@ -40,7 +74,7 @@ impl ArrayChunks { #[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[..self.init]) } + unsafe { MaybeUninit::slice_assume_init_ref(&self.guard.arr[..self.guard.init]) } } /// Returns the remainder of the elements yielded by the original @@ -49,7 +83,7 @@ impl ArrayChunks { #[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[..self.init]) } + unsafe { MaybeUninit::slice_assume_init_mut(&mut self.guard.arr[..self.guard.init]) } } } @@ -58,14 +92,72 @@ impl Iterator for ArrayChunks { type Item = [I::Item; N]; fn next(&mut self) -> Option { - while self.init < N { - self.buffer[self.init] = MaybeUninit::new(self.iter.next()?); - self.init += 1; + while self.guard.init < N { + self.guard.arr[self.guard.init] = MaybeUninit::new(self.iter.next()?); + self.guard.init += 1; } - self.init = 0; - // SAFETY: This is safe: `MaybeUninit` is guaranteed to have the same layout - // as `T` and the entire array has just been initialized. - unsafe { Some(mem::transmute_copy(&self.buffer)) } + // SAFETY: The entire array has just been initialized. + unsafe { Some(self.guard.assume_init_read()) } + } + + fn size_hint(&self) -> (usize, Option) { + 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 => self.iter.advance_by(usize::MAX).and(Err(usize::MAX)), + }; + res.map_err(|k| k / N) + } + + fn try_fold(&mut self, init: Acc, mut fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try, + { + let mut guard = mem::replace(&mut self.guard, Guard::new()); + let result = self.iter.try_fold(init, |mut acc, x| { + guard.arr[guard.init] = MaybeUninit::new(x); + guard.init += 1; + + if guard.init == N { + // SAFETY: The entire array has just been initialized. + let arr = unsafe { guard.assume_init_read() }; + acc = fold(acc, arr).branch()?; + } + ControlFlow::Continue(acc) + }); + + if guard.init != 0 { + self.guard = guard; + } + + match result { + ControlFlow::Continue(acc) => R::from_output(acc), + ControlFlow::Break(res) => R::from_residual(res), + } + } + + fn fold(self, init: Acc, mut fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + let mut guard = Guard::<_, N>::new(); + self.iter.fold(init, |mut acc, x| { + guard.arr[guard.init] = MaybeUninit::new(x); + guard.init += 1; + + if guard.init == N { + // SAFETY: The entire array has just been initialized. + let arr = unsafe { guard.assume_init_read() }; + acc = fold(acc, arr); + } + acc + }) } } diff --git a/library/core/tests/iter/adapters/array_chunks.rs b/library/core/tests/iter/adapters/array_chunks.rs index 67cb7ff44f5ff..63c0d2ca31c74 100644 --- a/library/core/tests/iter/adapters/array_chunks.rs +++ b/library/core/tests/iter/adapters/array_chunks.rs @@ -1,12 +1,18 @@ use core::cell::Cell; use core::iter::*; -#[derive(Debug)] +#[derive(Debug, Clone)] struct DropBomb<'a> { dropped: bool, counter: &'a Cell, } +impl<'a> DropBomb<'a> { + fn new(counter: &'a Cell) -> Self { + Self { dropped: false, counter } + } +} + impl Drop for DropBomb<'_> { fn drop(&mut self) { if self.dropped { @@ -17,6 +23,34 @@ impl Drop for DropBomb<'_> { } } +#[test] +fn test_iterator_array_chunks_clone() { + let mut iter = (0..=10).array_chunks::<4>(); + let mut iter2 = iter.clone(); + for (x, y) in iter.by_ref().zip(iter2.by_ref()) { + assert_eq!(x, y); + } + assert_eq!(iter.remainder(), &[8, 9, 10]); + assert_eq!(iter2.remainder(), &[]); + assert_eq!(iter2.next(), None); + assert_eq!(iter2.remainder(), &[8, 9, 10]); + + let counter = Cell::new(0); + let mut iter = once(DropBomb::new(&counter)).cycle().take(11).array_chunks::<3>(); + let mut iter2 = iter.clone(); + for (i, (_x, _y)) in iter.by_ref().zip(iter2.by_ref()).enumerate() { + assert_eq!(counter.get(), i * 6); + } + assert_eq!(counter.get(), 18); + drop(iter); + assert_eq!(counter.get(), 21); + assert_eq!(iter2.remainder().len(), 0); + assert!(iter2.next().is_none()); + assert_eq!(iter2.remainder().len(), 2); + drop(iter2); + assert_eq!(counter.get(), 24); +} + #[test] fn test_iterator_array_chunks_remainder() { let mut iter = (0..=10).array_chunks::<4>(); @@ -36,8 +70,7 @@ fn test_iterator_array_chunks_remainder() { #[test] fn test_iterator_array_chunks_drop() { let counter = Cell::new(0); - let create = - |n| (0..n).map(|_| DropBomb { dropped: false, counter: &counter }).array_chunks::<3>(); + let create = |n| (0..n).map(|_| DropBomb::new(&counter)).array_chunks::<3>(); let iter = create(5); assert_eq!(counter.get(), 0); @@ -64,3 +97,64 @@ fn test_iterator_array_chunks_drop() { drop(iter); assert_eq!(counter.get(), 5); } + +#[test] +fn test_iterator_array_chunks_try_fold() { + let mut iter = (0..=10).array_chunks::<3>(); + let result = iter.try_fold(0, |acc, arr| { + assert_eq!(arr, [acc * 3, (acc * 3) + 1, (acc * 3) + 2]); + if acc == 2 { Err(acc) } else { Ok(acc + 1) } + }); + assert_eq!(result, Err(2)); + assert_eq!(iter.next(), None); + assert_eq!(iter.remainder(), &[9, 10]); + + let mut iter = (0..10).array_chunks::<2>(); + let result: Result<_, ()> = iter.try_fold(0, |acc, arr| { + assert_eq!(arr, [acc * 2, (acc * 2) + 1]); + Ok(acc + 1) + }); + assert_eq!(result, Ok(5)); + assert_eq!(iter.next(), None); + assert_eq!(iter.remainder(), &[]); + + let counter = Cell::new(0); + let mut iter = (0..=10).map(|_| DropBomb::new(&counter)).array_chunks::<3>(); + let result = iter.try_fold(0, |acc, _arr| { + assert_eq!(counter.get(), acc * 3); + if acc == 1 { Err(acc) } else { Ok(acc + 1) } + }); + assert_eq!(result, Err(1)); + assert_eq!(iter.remainder().len(), 0); + assert_eq!(counter.get(), 6); + drop(iter); + assert_eq!(counter.get(), 6); + + counter.set(0); + let mut iter = (0..=10).map(|_| DropBomb::new(&counter)).array_chunks::<3>(); + let result: Result<_, ()> = iter.try_fold(0, |acc, _arr| { + assert_eq!(counter.get(), acc * 3); + Ok(acc + 1) + }); + assert_eq!(result, Ok(3)); + assert_eq!(iter.remainder().len(), 2); + assert_eq!(counter.get(), 9); + drop(iter); + assert_eq!(counter.get(), 11); +} + +#[test] +fn test_iterator_array_chunks_fold() { + let result = (0..10).array_chunks::<3>().fold(0, |acc, arr| { + assert_eq!(arr, [acc * 3, (acc * 3) + 1, (acc * 3) + 2]); + acc + 1 + }); + assert_eq!(result, 3); + + let counter = Cell::new(0); + (0..10).map(|_| DropBomb::new(&counter)).array_chunks::<3>().fold(0, |acc, _arr| { + assert_eq!(counter.get(), acc * 3); + acc + 1 + }); + assert_eq!(counter.get(), 10); +} From bd9d277e0cc44ff40481f99c1ea674fea1022072 Mon Sep 17 00:00:00 2001 From: John Schug Date: Thu, 5 Aug 2021 18:39:51 -0700 Subject: [PATCH 4/8] Add benchmarks for Iterator::array_chunks --- library/core/benches/iter.rs | 41 ++++++++++++++++++++++++++++++++++++ library/core/benches/lib.rs | 2 ++ 2 files changed, 43 insertions(+) diff --git a/library/core/benches/iter.rs b/library/core/benches/iter.rs index 24257ba98785d..d6c9a2b166f75 100644 --- a/library/core/benches/iter.rs +++ b/library/core/benches/iter.rs @@ -367,3 +367,44 @@ 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_iter_array_chunks(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::()); + 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::()); + acc + }); +} + +#[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::()); + 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::()); + acc + }); +} diff --git a/library/core/benches/lib.rs b/library/core/benches/lib.rs index de4ef7949f344..5dbb9e559521c 100644 --- a/library/core/benches/lib.rs +++ b/library/core/benches/lib.rs @@ -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; From 2289b7727df47b0a275e5ff1cd9ee505bfec9f48 Mon Sep 17 00:00:00 2001 From: John Schug Date: Fri, 6 Aug 2021 00:43:08 -0700 Subject: [PATCH 5/8] Improve performance of Iterator::array_chunks --- library/core/benches/iter.rs | 18 +- .../core/src/iter/adapters/array_chunks.rs | 195 ++++++++++++------ 2 files changed, 145 insertions(+), 68 deletions(-) diff --git a/library/core/benches/iter.rs b/library/core/benches/iter.rs index d6c9a2b166f75..728cc29570a18 100644 --- a/library/core/benches/iter.rs +++ b/library/core/benches/iter.rs @@ -369,7 +369,17 @@ fn bench_lt(b: &mut Bencher) { } #[bench] -fn bench_iter_array_chunks(b: &mut Bencher) { +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::()); + 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); @@ -389,11 +399,11 @@ fn bench_iter_array_chunks_fold(b: &mut Bencher) { } #[bench] -fn bench_array_map(b: &mut Bencher) { +fn bench_iter_array_chunks_ref_fold(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::()); + let mut iter = (0i64..1000000).array_chunks::<100>().map(black_box); + for_each_fold(iter.by_ref(), |x| acc += x.iter().sum::()); acc }); } diff --git a/library/core/src/iter/adapters/array_chunks.rs b/library/core/src/iter/adapters/array_chunks.rs index 4e726579ad9b0..9bd1494675b08 100644 --- a/library/core/src/iter/adapters/array_chunks.rs +++ b/library/core/src/iter/adapters/array_chunks.rs @@ -4,51 +4,86 @@ use core::ops::{ControlFlow, Try}; use core::ptr; #[derive(Debug)] -struct Guard { - arr: [MaybeUninit; N], +struct Buffer { + array: [MaybeUninit; N], init: usize, } -impl Guard { +impl Buffer { fn new() -> Self { - Self { arr: MaybeUninit::uninit_array(), init: 0 } - } - - // SAFETY: The entire array must have been properly initialized. - unsafe fn assume_init_read(&mut self) -> [T; N] { - debug_assert_eq!(self.init, N); - // Reset the tracked initialization state. - self.init = 0; - // SAFETY: This is safe: `MaybeUninit` is guaranteed to have the same layout - // as `T` and per the function's contract the entire array has been initialized. - unsafe { mem::transmute_copy(&self.arr) } + Self { array: MaybeUninit::uninit_array(), init: 0 } } } -impl Clone for Guard { +impl Clone for Buffer { 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.arr[..self.init]) }; - MaybeUninit::write_slice_cloned(&mut new.arr[..self.init], src); + 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 Drop for Guard { +impl Drop for Buffer { fn drop(&mut self) { debug_assert!(self.init <= N); - let init_part = &mut self.arr[..self.init]; + 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(init_part)); + ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(initialized_part)); } } } +// FIXME: Combine with `Guard` in `collect_into_array`. +struct Guard { + ptr: *mut T, + init: usize, +} + +impl Drop for Guard { + 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 Guard { + fn new(array: &mut [MaybeUninit; N]) -> Self { + Self { ptr: MaybeUninit::slice_as_mut_ptr(array), init: 0 } + } + + fn with(buffer: &mut Buffer, f: F) -> R + where + F: FnOnce(&mut [MaybeUninit; 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`. /// @@ -60,12 +95,12 @@ impl Drop for Guard { #[derive(Debug, Clone)] pub struct ArrayChunks { iter: I, - guard: Guard, + buffer: Buffer, } impl ArrayChunks { pub(in crate::iter) fn new(iter: I) -> Self { - Self { iter, guard: Guard::new() } + Self { iter, buffer: Buffer::new() } } /// Returns the remainder of the elements yielded by the original @@ -74,7 +109,7 @@ impl ArrayChunks { #[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.guard.arr[..self.guard.init]) } + unsafe { MaybeUninit::slice_assume_init_ref(&self.buffer.array[..self.buffer.init]) } } /// Returns the remainder of the elements yielded by the original @@ -83,7 +118,7 @@ impl ArrayChunks { #[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.guard.arr[..self.guard.init]) } + unsafe { MaybeUninit::slice_assume_init_mut(&mut self.buffer.array[..self.buffer.init]) } } } @@ -92,12 +127,21 @@ impl Iterator for ArrayChunks { type Item = [I::Item; N]; fn next(&mut self) -> Option { - while self.guard.init < N { - self.guard.arr[self.guard.init] = MaybeUninit::new(self.iter.next()?); - self.guard.init += 1; - } - // SAFETY: The entire array has just been initialized. - unsafe { Some(self.guard.assume_init_read()) } + 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) { @@ -108,55 +152,78 @@ impl Iterator for ArrayChunks { fn advance_by(&mut self, n: usize) -> Result<(), usize> { let res = match n.checked_mul(N) { Some(n) => self.iter.advance_by(n), - None => self.iter.advance_by(usize::MAX).and(Err(usize::MAX)), + None => { + let n = (usize::MAX / N) * N; + self.iter.advance_by(n).and(Err(n)) + } }; res.map_err(|k| k / N) } - fn try_fold(&mut self, init: Acc, mut fold: Fold) -> R + fn try_fold(&mut self, acc: Acc, mut fold: Fold) -> R where Self: Sized, Fold: FnMut(Acc, Self::Item) -> R, R: Try, { - let mut guard = mem::replace(&mut self.guard, Guard::new()); - let result = self.iter.try_fold(init, |mut acc, x| { - guard.arr[guard.init] = MaybeUninit::new(x); - guard.init += 1; - - if guard.init == N { - // SAFETY: The entire array has just been initialized. - let arr = unsafe { guard.assume_init_read() }; - acc = fold(acc, arr).branch()?; + 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), } - ControlFlow::Continue(acc) - }); - - if guard.init != 0 { - self.guard = guard; - } - - match result { - ControlFlow::Continue(acc) => R::from_output(acc), - ControlFlow::Break(res) => R::from_residual(res), - } + }) } - fn fold(self, init: Acc, mut fold: Fold) -> Acc + fn fold(self, acc: Acc, mut fold: Fold) -> Acc where Fold: FnMut(Acc, Self::Item) -> Acc, { - let mut guard = Guard::<_, N>::new(); - self.iter.fold(init, |mut acc, x| { - guard.arr[guard.init] = MaybeUninit::new(x); - guard.init += 1; - - if guard.init == N { - // SAFETY: The entire array has just been initialized. - let arr = unsafe { guard.assume_init_read() }; - acc = fold(acc, arr); - } - 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 + }) }) } } From 6a15bc6dfb679b66cad0a34ad931be4bf2c42a9c Mon Sep 17 00:00:00 2001 From: John Schug Date: Fri, 6 Aug 2021 12:02:36 -0700 Subject: [PATCH 6/8] Add DoubleEndedIterator impl for ArrayChunks --- .../core/src/iter/adapters/array_chunks.rs | 118 +++++++++++++++--- .../core/tests/iter/adapters/array_chunks.rs | 25 ++++ 2 files changed, 125 insertions(+), 18 deletions(-) diff --git a/library/core/src/iter/adapters/array_chunks.rs b/library/core/src/iter/adapters/array_chunks.rs index 9bd1494675b08..00911999e7a97 100644 --- a/library/core/src/iter/adapters/array_chunks.rs +++ b/library/core/src/iter/adapters/array_chunks.rs @@ -1,6 +1,6 @@ use core::iter::FusedIterator; use core::mem::{self, MaybeUninit}; -use core::ops::{ControlFlow, Try}; +use core::ops::Try; use core::ptr; #[derive(Debug)] @@ -149,18 +149,97 @@ impl Iterator for ArrayChunks { (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)) + fn try_fold(&mut self, acc: Acc, mut fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try, + { + let iter = &mut self.iter; + Guard::with(&mut self.buffer, |array, init| { + 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)?; + } + R::from_output(acc) + }) + }) + } + + fn 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 DoubleEndedIterator for ArrayChunks { + fn next_back(&mut self) -> Option { + let iter = &mut self.iter; + let result = Guard::with(&mut self.buffer, |array, init| { + for slot in &mut array[*init..] { + slot.write(iter.next_back()?); + *init += 1; } - }; - res.map_err(|k| k / N) + *init = 0; + array.reverse(); + // SAFETY: The entire array has just been initialized. + unsafe { + Some(MaybeUninit::array_assume_init(mem::replace( + array, + MaybeUninit::uninit_array(), + ))) + } + }); + + if self.buffer.init > 0 { + (&mut self.buffer.array[..self.buffer.init]).reverse(); + } + + result } - fn try_fold(&mut self, acc: Acc, mut fold: Fold) -> R + fn try_rfold(&mut self, acc: Acc, mut fold: Fold) -> R where Self: Sized, Fold: FnMut(Acc, Self::Item) -> R, @@ -168,7 +247,7 @@ impl Iterator for ArrayChunks { { let iter = &mut self.iter; Guard::with(&mut self.buffer, |array, init| { - let result = iter.try_fold(acc, |mut acc, x| { + let result = iter.try_rfold(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 { @@ -178,6 +257,7 @@ impl Iterator for ArrayChunks { if *init == N { *init = 0; + array.reverse(); // SAFETY: The entire array has just been initialized. let array = unsafe { MaybeUninit::array_assume_init(mem::replace( @@ -185,25 +265,26 @@ impl Iterator for ArrayChunks { MaybeUninit::uninit_array(), )) }; - acc = fold(acc, array).branch()?; + acc = fold(acc, array)?; } - ControlFlow::Continue(acc) + R::from_output(acc) }); - match result { - ControlFlow::Continue(acc) => R::from_output(acc), - ControlFlow::Break(res) => R::from_residual(res), + if *init > 0 { + (&mut array[..*init]).reverse(); } + + result }) } - fn fold(self, acc: Acc, mut fold: Fold) -> Acc + fn rfold(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| { + iter.rfold(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 { @@ -213,6 +294,7 @@ impl Iterator for ArrayChunks { if *init == N { *init = 0; + array.reverse(); // SAFETY: The entire array has just been initialized. let array = unsafe { MaybeUninit::array_assume_init(mem::replace( diff --git a/library/core/tests/iter/adapters/array_chunks.rs b/library/core/tests/iter/adapters/array_chunks.rs index 63c0d2ca31c74..1568838b9d84f 100644 --- a/library/core/tests/iter/adapters/array_chunks.rs +++ b/library/core/tests/iter/adapters/array_chunks.rs @@ -158,3 +158,28 @@ fn test_iterator_array_chunks_fold() { }); assert_eq!(counter.get(), 10); } + +#[test] +fn test_iterator_array_chunks_rev() { + let mut iter = (0..=10).array_chunks::<4>(); + assert_eq!(iter.next_back(), Some([7, 8, 9, 10])); + assert_eq!(iter.next_back(), Some([3, 4, 5, 6])); + assert_eq!(iter.next_back(), None); + assert_eq!(iter.remainder(), &[0, 1, 2]); +} + +#[test] +fn test_iterator_array_chunks_rfold() { + let result = (0..10).array_chunks::<3>().rfold(0, |acc, arr| { + assert_eq!(arr, [7 - (acc * 3), 8 - (acc * 3), 9 - (acc * 3)]); + acc + 1 + }); + assert_eq!(result, 3); + + let counter = Cell::new(0); + (0..10).map(|_| DropBomb::new(&counter)).array_chunks::<3>().rfold(0, |acc, _arr| { + assert_eq!(counter.get(), acc * 3); + acc + 1 + }); + assert_eq!(counter.get(), 10); +} From c14ba10f9c10f9314e32a9282cc9c97ac47501c0 Mon Sep 17 00:00:00 2001 From: John Schug Date: Fri, 6 Aug 2021 13:26:23 -0700 Subject: [PATCH 7/8] Make ArrayChunks unconditionally fused --- library/core/benches/iter.rs | 21 +- library/core/src/array/mod.rs | 66 +-- .../core/src/iter/adapters/array_chunks.rs | 381 ++++++++++-------- .../core/tests/iter/adapters/array_chunks.rs | 49 +++ 4 files changed, 302 insertions(+), 215 deletions(-) diff --git a/library/core/benches/iter.rs b/library/core/benches/iter.rs index 728cc29570a18..aed3b2a2be9d7 100644 --- a/library/core/benches/iter.rs +++ b/library/core/benches/iter.rs @@ -369,11 +369,11 @@ fn bench_lt(b: &mut Bencher) { } #[bench] -fn bench_array_map(b: &mut Bencher) { +fn bench_iter_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::()); + for_each_fold(iter, |x: [i64; 100]| acc += x.iter().sum::()); acc }); } @@ -383,7 +383,7 @@ 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::()); + for_each_loop(iter, |x: [i64; 100]| acc += x.iter().sum::()); acc }); } @@ -393,28 +393,27 @@ 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::()); + for_each_fold(iter, |x: [i64; 100]| acc += x.iter().sum::()); acc }); } #[bench] -fn bench_iter_array_chunks_ref_fold(b: &mut Bencher) { +fn bench_iter_array_chunks_rev_loop(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::()); + let iter = (0i64..1000000).array_chunks::<100>().map(black_box); + for_each_loop(iter.rev(), |x: [i64; 100]| acc += x.iter().sum::()); acc }); } #[bench] -fn bench_slice_array_chunks(b: &mut Bencher) { - let vec: Vec<_> = (0i64..1000000).collect(); +fn bench_iter_array_chunks_rfold(b: &mut Bencher) { b.iter(|| { let mut acc = 0; - let iter = vec.array_chunks::<100>().map(black_box); - for_each_loop(iter, |x| acc += x.iter().sum::()); + let iter = (0i64..1000000).array_chunks::<100>().map(black_box); + for_each_fold(iter.rev(), |x: [i64; 100]| acc += x.iter().sum::()); acc }); } diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index 3bc9f71375cb8..f86c6a7f1786a 100644 --- a/library/core/src/array/mod.rs +++ b/library/core/src/array/mod.rs @@ -437,6 +437,38 @@ impl [T; N] { } } +pub(crate) struct Guard { + ptr: *mut T, + pub(crate) init: usize, +} + +impl Drop for Guard { + fn drop(&mut self) { + debug_assert!(self.init <= N); + + let init_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(init_part); + } + } +} + +impl Guard { + /// Creates a `Guard` object that will drop the initialized portion + /// of the provided array when dropped. + /// + /// # Safety + /// + /// When the returned guard is dropped, the array must still be valid, + /// `init` must not be greater `N`, and the first `init` elements of + /// the array must be properly initialized. + pub(crate) unsafe fn new(array: &mut [MaybeUninit; N]) -> Self { + Self { ptr: MaybeUninit::slice_as_mut_ptr(array), init: 0 } + } +} + /// Pulls `N` items from `iter` and returns them as an array. If the iterator /// yields fewer than `N` items, this function exhibits undefined behavior. /// @@ -484,39 +516,23 @@ where return unsafe { Some(mem::zeroed()) }; } - struct Guard { - ptr: *mut T, - initialized: usize, - } - - impl Drop for Guard { - fn drop(&mut self) { - debug_assert!(self.initialized <= N); - - let initialized_part = crate::ptr::slice_from_raw_parts_mut(self.ptr, self.initialized); - - // SAFETY: this raw slice will contain only initialized objects. - unsafe { - crate::ptr::drop_in_place(initialized_part); - } - } - } - let mut array = MaybeUninit::uninit_array::(); - let mut guard: Guard<_, N> = - Guard { ptr: MaybeUninit::slice_as_mut_ptr(&mut array), initialized: 0 }; + // SAFETY: `guard` is always either forgotten or dropped before the array + // is moved/dropped and `guard.init` properly tracks the initialized + // members of the array. + let mut guard = unsafe { Guard::new(&mut array) }; while let Some(item) = iter.next() { - // SAFETY: `guard.initialized` starts at 0, is increased by one in the + // SAFETY: `guard.init` starts at 0, is increased by one in the // loop and the loop is aborted once it reaches N (which is // `array.len()`). unsafe { - array.get_unchecked_mut(guard.initialized).write(item); + array.get_unchecked_mut(guard.init).write(item); } - guard.initialized += 1; + guard.init += 1; // Check if the whole array was initialized. - if guard.initialized == N { + if guard.init == N { mem::forget(guard); // SAFETY: the condition above asserts that all elements are @@ -527,7 +543,7 @@ where } // This is only reached if the iterator is exhausted before - // `guard.initialized` reaches `N`. Also note that `guard` is dropped here, + // `guard.init` reaches `N`. Also note that `guard` is dropped here, // dropping all already initialized elements. None } diff --git a/library/core/src/iter/adapters/array_chunks.rs b/library/core/src/iter/adapters/array_chunks.rs index 00911999e7a97..f8b7c8557de3f 100644 --- a/library/core/src/iter/adapters/array_chunks.rs +++ b/library/core/src/iter/adapters/array_chunks.rs @@ -1,6 +1,7 @@ +use core::array::Guard; use core::iter::FusedIterator; use core::mem::{self, MaybeUninit}; -use core::ops::Try; +use core::ops::{ControlFlow, Try}; use core::ptr; #[derive(Debug)] @@ -40,50 +41,6 @@ impl Drop for Buffer { } } -// FIXME: Combine with `Guard` in `collect_into_array`. -struct Guard { - ptr: *mut T, - init: usize, -} - -impl Drop for Guard { - 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 Guard { - fn new(array: &mut [MaybeUninit; N]) -> Self { - Self { ptr: MaybeUninit::slice_as_mut_ptr(array), init: 0 } - } - - fn with(buffer: &mut Buffer, f: F) -> R - where - F: FnOnce(&mut [MaybeUninit; 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`. /// @@ -96,11 +53,12 @@ impl Guard { pub struct ArrayChunks { iter: I, buffer: Buffer, + done: bool, } impl ArrayChunks { pub(in crate::iter) fn new(iter: I) -> Self { - Self { iter, buffer: Buffer::new() } + Self { iter, buffer: Buffer::new(), done: false } } /// Returns the remainder of the elements yielded by the original @@ -127,87 +85,126 @@ impl Iterator for ArrayChunks { type Item = [I::Item; N]; fn next(&mut self) -> Option { - 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(), - ))) - } - }) + if self.done { + return None; + } + + let mut array = MaybeUninit::uninit_array(); + // SAFETY: `guard` is always either forgotten or dropped before the array + // is moved/dropped and `guard.init` properly tracks the initialized + // members of the array. + let mut guard = unsafe { Guard::new(&mut array) }; + for slot in &mut array { + let next = match self.iter.next() { + Some(n) => n, + None => { + self.done = true; + if guard.init > 0 { + self.buffer.init = guard.init; + mem::forget(guard); + self.buffer.array = array; + } + return None; + } + }; + slot.write(next); + guard.init += 1; + } + mem::forget(guard); + // SAFETY: The entire array has just been initialized. + unsafe { Some(MaybeUninit::array_assume_init(array)) } } fn size_hint(&self) -> (usize, Option) { + if self.done { + return (0, Some(0)); + } let (lower, upper) = self.iter.size_hint(); (lower / N, upper.map(|x| x / N)) } fn try_fold(&mut self, acc: Acc, mut fold: Fold) -> R where - Self: Sized, Fold: FnMut(Acc, Self::Item) -> R, R: Try, { - let iter = &mut self.iter; - Guard::with(&mut self.buffer, |array, init| { - 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 self.done { + return R::from_output(acc); + } - 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)?; + let mut array = MaybeUninit::uninit_array(); + // SAFETY: `guard` is always either forgotten or dropped before the array + // is moved/dropped and `guard.init` properly tracks the initialized + // members of the array. + let mut guard = unsafe { Guard::new(&mut array) }; + let result = self.iter.try_fold(acc, |mut acc, x| { + // SAFETY: `guard.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(guard.init).write(x); + } + guard.init += 1; + + if guard.init == N { + guard.init = 0; + // SAFETY: The entire array has just been initialized. + let array = unsafe { + MaybeUninit::array_assume_init(mem::replace( + &mut array, + MaybeUninit::uninit_array(), + )) + }; + acc = fold(acc, array)?; + } + R::from_output(acc) + }); + match result.branch() { + ControlFlow::Continue(x) => { + self.done = true; + if guard.init > 0 { + self.buffer.init = guard.init; + mem::forget(guard); + self.buffer.array = array; } - R::from_output(acc) - }) - }) + R::from_output(x) + } + ControlFlow::Break(x) => R::from_residual(x), + } } fn 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 self.done { + return acc; + } - 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 - }) + let mut array = MaybeUninit::uninit_array(); + // SAFETY: `guard` is always either forgotten or dropped before the array + // is moved/dropped and `guard.init` properly tracks the initialized + // members of the array. + let mut guard = unsafe { Guard::new(&mut array) }; + self.iter.fold(acc, |mut acc, x| { + // SAFETY: `guard.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(guard.init).write(x); + } + guard.init += 1; + + if guard.init == N { + guard.init = 0; + // SAFETY: The entire array has just been initialized. + let array = unsafe { + MaybeUninit::array_assume_init(mem::replace( + &mut array, + MaybeUninit::uninit_array(), + )) + }; + acc = fold(acc, array); + } + acc }) } } @@ -215,100 +212,126 @@ impl Iterator for ArrayChunks { #[unstable(feature = "iter_array_chunks", issue = "none")] impl DoubleEndedIterator for ArrayChunks { fn next_back(&mut self) -> Option { - let iter = &mut self.iter; - let result = Guard::with(&mut self.buffer, |array, init| { - for slot in &mut array[*init..] { - slot.write(iter.next_back()?); - *init += 1; - } - *init = 0; - array.reverse(); - // SAFETY: The entire array has just been initialized. - unsafe { - Some(MaybeUninit::array_assume_init(mem::replace( - array, - MaybeUninit::uninit_array(), - ))) - } - }); - - if self.buffer.init > 0 { - (&mut self.buffer.array[..self.buffer.init]).reverse(); + if self.done { + return None; } - result + let mut array = MaybeUninit::uninit_array(); + // SAFETY: `guard` is always either forgotten or dropped before the array + // is moved/dropped and `guard.init` properly tracks the initialized + // members of the array. + let mut guard = unsafe { Guard::new(&mut array) }; + for slot in &mut array { + let next = match self.iter.next_back() { + Some(n) => n, + None => { + self.done = true; + if guard.init > 0 { + (&mut array[..guard.init]).reverse(); + self.buffer.init = guard.init; + mem::forget(guard); + self.buffer.array = array; + } + return None; + } + }; + slot.write(next); + guard.init += 1; + } + guard.init = 0; + array.reverse(); + // SAFETY: The entire array has just been initialized. + unsafe { Some(MaybeUninit::array_assume_init(array)) } } fn try_rfold(&mut self, acc: Acc, mut fold: Fold) -> R where - Self: Sized, Fold: FnMut(Acc, Self::Item) -> R, R: Try, { - let iter = &mut self.iter; - Guard::with(&mut self.buffer, |array, init| { - let result = iter.try_rfold(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 self.done { + return R::from_output(acc); + } - if *init == N { - *init = 0; - array.reverse(); - // 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)?; + let mut array = MaybeUninit::uninit_array(); + // SAFETY: `guard` is always either forgotten or dropped before the array + // is moved/dropped and `guard.init` properly tracks the initialized + // members of the array. + let mut guard = unsafe { Guard::new(&mut array) }; + let result = self.iter.try_rfold(acc, |mut acc, x| { + // SAFETY: `guard.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(guard.init).write(x); + } + guard.init += 1; + + if guard.init == N { + guard.init = 0; + array.reverse(); + // SAFETY: The entire array has just been initialized. + let array = unsafe { + MaybeUninit::array_assume_init(mem::replace( + &mut array, + MaybeUninit::uninit_array(), + )) + }; + acc = fold(acc, array)?; + } + R::from_output(acc) + }); + match result.branch() { + ControlFlow::Continue(x) => { + self.done = true; + if guard.init > 0 { + (&mut array[..guard.init]).reverse(); + self.buffer.init = guard.init; + mem::forget(guard); + self.buffer.array = array; } - R::from_output(acc) - }); - - if *init > 0 { - (&mut array[..*init]).reverse(); + R::from_output(x) } - - result - }) + ControlFlow::Break(x) => R::from_residual(x), + } } fn rfold(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.rfold(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 self.done { + return acc; + } - if *init == N { - *init = 0; - array.reverse(); - // 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 - }) + let mut array = MaybeUninit::uninit_array(); + // SAFETY: `guard` is always either forgotten or dropped before the array + // is moved/dropped and `guard.init` properly tracks the initialized + // members of the array. + let mut guard = unsafe { Guard::new(&mut array) }; + self.iter.rfold(acc, |mut acc, x| { + // SAFETY: `guard.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(guard.init).write(x); + } + guard.init += 1; + + if guard.init == N { + guard.init = 0; + array.reverse(); + // SAFETY: The entire array has just been initialized. + let array = unsafe { + MaybeUninit::array_assume_init(mem::replace( + &mut array, + MaybeUninit::uninit_array(), + )) + }; + acc = fold(acc, array); + } + acc }) } } #[unstable(feature = "iter_array_chunks", issue = "none")] -impl FusedIterator for ArrayChunks {} +impl FusedIterator for ArrayChunks {} diff --git a/library/core/tests/iter/adapters/array_chunks.rs b/library/core/tests/iter/adapters/array_chunks.rs index 1568838b9d84f..77531e4bea9a5 100644 --- a/library/core/tests/iter/adapters/array_chunks.rs +++ b/library/core/tests/iter/adapters/array_chunks.rs @@ -63,6 +63,7 @@ fn test_iterator_array_chunks_remainder() { assert_eq!(iter.remainder(), &[]); assert_eq!(iter.remainder_mut(), &[]); assert_eq!(iter.next(), None); + assert_eq!(iter.next(), None); assert_eq!(iter.remainder(), &[8, 9, 10]); assert_eq!(iter.remainder_mut(), &[8, 9, 10]); } @@ -106,6 +107,7 @@ fn test_iterator_array_chunks_try_fold() { if acc == 2 { Err(acc) } else { Ok(acc + 1) } }); assert_eq!(result, Err(2)); + assert_eq!(iter.remainder(), &[]); assert_eq!(iter.next(), None); assert_eq!(iter.remainder(), &[9, 10]); @@ -165,9 +167,56 @@ fn test_iterator_array_chunks_rev() { assert_eq!(iter.next_back(), Some([7, 8, 9, 10])); assert_eq!(iter.next_back(), Some([3, 4, 5, 6])); assert_eq!(iter.next_back(), None); + assert_eq!(iter.next_back(), None); assert_eq!(iter.remainder(), &[0, 1, 2]); } +#[test] +fn test_iterator_array_chunks_try_rfold() { + let mut iter = (0..=10).array_chunks::<3>(); + let result = iter.try_rfold(0, |acc, arr| { + assert_eq!(arr, [8 - (acc * 3), 9 - (acc * 3), 10 - (acc * 3)]); + if acc == 2 { Err(acc) } else { Ok(acc + 1) } + }); + assert_eq!(result, Err(2)); + assert_eq!(iter.remainder(), &[]); + assert_eq!(iter.next(), None); + assert_eq!(iter.remainder(), &[0, 1]); + + let mut iter = (0..10).array_chunks::<2>(); + let result: Result<_, ()> = iter.try_rfold(0, |acc, arr| { + assert_eq!(arr, [8 - (acc * 2), 9 - (acc * 2)]); + Ok(acc + 1) + }); + assert_eq!(result, Ok(5)); + assert_eq!(iter.next(), None); + assert_eq!(iter.remainder(), &[]); + + let counter = Cell::new(0); + let mut iter = (0..=10).map(|_| DropBomb::new(&counter)).array_chunks::<3>(); + let result = iter.try_rfold(0, |acc, _arr| { + assert_eq!(counter.get(), acc * 3); + if acc == 1 { Err(acc) } else { Ok(acc + 1) } + }); + assert_eq!(result, Err(1)); + assert_eq!(iter.remainder().len(), 0); + assert_eq!(counter.get(), 6); + drop(iter); + assert_eq!(counter.get(), 6); + + counter.set(0); + let mut iter = (0..=10).map(|_| DropBomb::new(&counter)).array_chunks::<3>(); + let result: Result<_, ()> = iter.try_rfold(0, |acc, _arr| { + assert_eq!(counter.get(), acc * 3); + Ok(acc + 1) + }); + assert_eq!(result, Ok(3)); + assert_eq!(iter.remainder().len(), 2); + assert_eq!(counter.get(), 9); + drop(iter); + assert_eq!(counter.get(), 11); +} + #[test] fn test_iterator_array_chunks_rfold() { let result = (0..10).array_chunks::<3>().rfold(0, |acc, arr| { From 85edaca03cd17047183a15902e424cd7226fd5a4 Mon Sep 17 00:00:00 2001 From: John Schug Date: Fri, 6 Aug 2021 14:21:46 -0700 Subject: [PATCH 8/8] Add ExactSizeIterator impl for ArrayChunks --- library/core/src/iter/adapters/array_chunks.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/library/core/src/iter/adapters/array_chunks.rs b/library/core/src/iter/adapters/array_chunks.rs index f8b7c8557de3f..d5974b9c60c73 100644 --- a/library/core/src/iter/adapters/array_chunks.rs +++ b/library/core/src/iter/adapters/array_chunks.rs @@ -333,5 +333,8 @@ impl DoubleEndedIterator for ArrayChunks } } +#[unstable(feature = "iter_array_chunks", issue = "none")] +impl ExactSizeIterator for ArrayChunks {} + #[unstable(feature = "iter_array_chunks", issue = "none")] impl FusedIterator for ArrayChunks {}