From b00666ed09f186d9891be463355d6f8c984dcd9e Mon Sep 17 00:00:00 2001 From: The 8472 Date: Sun, 23 Oct 2022 19:16:49 +0200 Subject: [PATCH 1/5] add benchmark for iter::ArrayChunks::fold specialization This also updates the existing iter::Copied::next_chunk benchmark so that the thing it benches doesn't get masked by the ArrayChunks specialization --- library/core/benches/iter.rs | 23 +++++++++++++++++++++-- library/core/benches/lib.rs | 1 + 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/library/core/benches/iter.rs b/library/core/benches/iter.rs index 38887f29af153..9193c79bee875 100644 --- a/library/core/benches/iter.rs +++ b/library/core/benches/iter.rs @@ -1,3 +1,4 @@ +use core::borrow::Borrow; use core::iter::*; use core::mem; use core::num::Wrapping; @@ -403,13 +404,31 @@ fn bench_trusted_random_access_adapters(b: &mut Bencher) { /// Exercises the iter::Copied specialization for slice::Iter #[bench] -fn bench_copied_array_chunks(b: &mut Bencher) { +fn bench_copied_chunks(b: &mut Bencher) { + let v = vec![1u8; 1024]; + + b.iter(|| { + let mut iter = black_box(&v).iter().copied(); + let mut acc = Wrapping(0); + // This uses a while-let loop to side-step the TRA specialization in ArrayChunks + while let Ok(chunk) = iter.next_chunk::<{ mem::size_of::() }>() { + let d = u64::from_ne_bytes(chunk); + acc += Wrapping(d.rotate_left(7).wrapping_add(1)); + } + acc + }) +} + +/// Exercises the TrustedRandomAccess specialization in ArrayChunks +#[bench] +fn bench_trusted_random_access_chunks(b: &mut Bencher) { let v = vec![1u8; 1024]; b.iter(|| { black_box(&v) .iter() - .copied() + // this shows that we're not relying on the slice::Iter specialization in Copied + .map(|b| *b.borrow()) .array_chunks::<{ mem::size_of::() }>() .map(|ary| { let d = u64::from_ne_bytes(ary); diff --git a/library/core/benches/lib.rs b/library/core/benches/lib.rs index 1e462e3fc3f8c..cf5a7cada93ce 100644 --- a/library/core/benches/lib.rs +++ b/library/core/benches/lib.rs @@ -5,6 +5,7 @@ #![feature(test)] #![feature(trusted_random_access)] #![feature(iter_array_chunks)] +#![feature(iter_next_chunk)] extern crate test; From eb3f001d3739e5e9f35e1a4b4e600889cf9980c6 Mon Sep 17 00:00:00 2001 From: The 8472 Date: Sun, 23 Oct 2022 19:17:41 +0200 Subject: [PATCH 2/5] make the array initialization guard available to other modules --- library/core/src/array/mod.rs | 36 +++++++++++++++++------------------ 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index eae0e1c761866..24e5d0fba21bf 100644 --- a/library/core/src/array/mod.rs +++ b/library/core/src/array/mod.rs @@ -865,24 +865,6 @@ where return Ok(Try::from_output(unsafe { mem::zeroed() })); } - struct Guard<'a, T, const N: usize> { - array_mut: &'a mut [MaybeUninit; N], - initialized: usize, - } - - impl Drop for Guard<'_, T, N> { - fn drop(&mut self) { - debug_assert!(self.initialized <= N); - - // SAFETY: this slice will contain only initialized objects. - unsafe { - crate::ptr::drop_in_place(MaybeUninit::slice_assume_init_mut( - &mut self.array_mut.get_unchecked_mut(..self.initialized), - )); - } - } - } - let mut array = MaybeUninit::uninit_array::(); let mut guard = Guard { array_mut: &mut array, initialized: 0 }; @@ -920,6 +902,24 @@ where Ok(Try::from_output(output)) } +pub(crate) struct Guard<'a, T, const N: usize> { + pub array_mut: &'a mut [MaybeUninit; N], + pub initialized: usize, +} + +impl Drop for Guard<'_, T, N> { + fn drop(&mut self) { + debug_assert!(self.initialized <= N); + + // SAFETY: this slice will contain only initialized objects. + unsafe { + crate::ptr::drop_in_place(MaybeUninit::slice_assume_init_mut( + &mut self.array_mut.get_unchecked_mut(..self.initialized), + )); + } + } +} + /// Returns the next chunk of `N` items from the iterator or errors with an /// iterator over the remainder. Used for `Iterator::next_chunk`. #[inline] From cfcce8e684c5e1bb2f9a74e55debf801ef27706f Mon Sep 17 00:00:00 2001 From: The 8472 Date: Sun, 23 Oct 2022 19:19:37 +0200 Subject: [PATCH 3/5] specialize iter::ArrayChunks::fold for TrustedRandomAccess iters This is fairly safe use of TRA since it consumes the iterator so no struct in an unsafe state will be left exposed to user code --- .../core/src/iter/adapters/array_chunks.rs | 89 ++++++++++++++++++- 1 file changed, 86 insertions(+), 3 deletions(-) diff --git a/library/core/src/iter/adapters/array_chunks.rs b/library/core/src/iter/adapters/array_chunks.rs index d4fb886101fe7..3f0fad4ed336f 100644 --- a/library/core/src/iter/adapters/array_chunks.rs +++ b/library/core/src/iter/adapters/array_chunks.rs @@ -1,6 +1,8 @@ use crate::array; -use crate::iter::{ByRefSized, FusedIterator, Iterator}; -use crate::ops::{ControlFlow, Try}; +use crate::const_closure::ConstFnMutClosure; +use crate::iter::{ByRefSized, FusedIterator, Iterator, TrustedRandomAccessNoCoerce}; +use crate::mem::{self, MaybeUninit}; +use crate::ops::{ControlFlow, NeverShortCircuit, Try}; /// An iterator over `N` elements of the iterator at a time. /// @@ -82,7 +84,13 @@ where } } - impl_fold_via_try_fold! { fold -> try_fold } + fn fold(self, init: B, f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B, + { + ::fold(self, init, f) + } } #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")] @@ -168,3 +176,78 @@ where self.iter.len() < N } } + +trait SpecFold: Iterator { + fn fold(self, init: B, f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B; +} + +impl SpecFold for ArrayChunks +where + I: Iterator, +{ + #[inline] + default fn fold(mut self, init: B, mut f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B, + { + let fold = ConstFnMutClosure::new(&mut f, NeverShortCircuit::wrap_mut_2_imp); + self.try_fold(init, fold).0 + } +} + +impl SpecFold for ArrayChunks +where + I: Iterator + TrustedRandomAccessNoCoerce, +{ + #[inline] + fn fold(mut self, init: B, mut f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B, + { + if self.remainder.is_some() { + return init; + } + + let mut accum = init; + let inner_len = self.iter.size(); + let mut i = 0; + // Use a while loop because (0..len).step_by(N) doesn't optimize well. + while inner_len - i >= N { + let mut chunk = MaybeUninit::uninit_array(); + let mut guard = array::Guard { array_mut: &mut chunk, initialized: 0 }; + for j in 0..N { + // SAFETY: The method consumes the iterator and the loop condition ensures that + // all accesses are in bounds and only happen once. + guard.array_mut[j].write(unsafe { self.iter.__iterator_get_unchecked(i + j) }); + guard.initialized = j + 1; + } + mem::forget(guard); + // SAFETY: The loop above initialized all elements + let chunk = unsafe { MaybeUninit::array_assume_init(chunk) }; + accum = f(accum, chunk); + i += N; + } + + let remainder = inner_len % N; + + let mut tail = MaybeUninit::uninit_array(); + let mut guard = array::Guard { array_mut: &mut tail, initialized: 0 }; + for i in 0..remainder { + // SAFETY: the remainder was not visited by the previous loop, so we're still only + // accessing each element once + let val = unsafe { self.iter.__iterator_get_unchecked(inner_len - remainder + i) }; + guard.array_mut[i].write(val); + guard.initialized = i + 1; + } + mem::forget(guard); + // SAFETY: the loop above initialized elements up to the `remainder` index + self.remainder = Some(unsafe { array::IntoIter::new_unchecked(tail, 0..remainder) }); + + accum + } +} From 43c353fff73a1ee38d0ec29042e42bdeea05c191 Mon Sep 17 00:00:00 2001 From: The 8472 Date: Sun, 23 Oct 2022 21:29:15 +0200 Subject: [PATCH 4/5] simplification: do not process the ArrayChunks remainder in fold() --- .../core/src/iter/adapters/array_chunks.rs | 20 ++----------------- .../core/tests/iter/adapters/array_chunks.rs | 3 ++- 2 files changed, 4 insertions(+), 19 deletions(-) diff --git a/library/core/src/iter/adapters/array_chunks.rs b/library/core/src/iter/adapters/array_chunks.rs index 3f0fad4ed336f..9bd135007b4a6 100644 --- a/library/core/src/iter/adapters/array_chunks.rs +++ b/library/core/src/iter/adapters/array_chunks.rs @@ -209,10 +209,6 @@ where Self: Sized, F: FnMut(B, Self::Item) -> B, { - if self.remainder.is_some() { - return init; - } - let mut accum = init; let inner_len = self.iter.size(); let mut i = 0; @@ -233,20 +229,8 @@ where i += N; } - let remainder = inner_len % N; - - let mut tail = MaybeUninit::uninit_array(); - let mut guard = array::Guard { array_mut: &mut tail, initialized: 0 }; - for i in 0..remainder { - // SAFETY: the remainder was not visited by the previous loop, so we're still only - // accessing each element once - let val = unsafe { self.iter.__iterator_get_unchecked(inner_len - remainder + i) }; - guard.array_mut[i].write(val); - guard.initialized = i + 1; - } - mem::forget(guard); - // SAFETY: the loop above initialized elements up to the `remainder` index - self.remainder = Some(unsafe { array::IntoIter::new_unchecked(tail, 0..remainder) }); + // unlike try_fold this method does not need to take care of the remainder + // since `self` will be dropped accum } diff --git a/library/core/tests/iter/adapters/array_chunks.rs b/library/core/tests/iter/adapters/array_chunks.rs index 4e9d89e1e580f..ef4a7e53bdd33 100644 --- a/library/core/tests/iter/adapters/array_chunks.rs +++ b/library/core/tests/iter/adapters/array_chunks.rs @@ -139,7 +139,8 @@ fn test_iterator_array_chunks_fold() { 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); + // fold impls may or may not process the remainder + assert!(count.get() <= 10 && count.get() >= 9); } #[test] From 3925fc0c8e0b4ad9be48fbc1136281841d8985fa Mon Sep 17 00:00:00 2001 From: The 8472 Date: Tue, 8 Nov 2022 00:13:26 +0100 Subject: [PATCH 5/5] document and improve array Guard type The type is unsafe and now exposed to the whole crate. Document it properly and add an unsafe method so the caller can make it visible that something unsafe is happening. --- library/core/src/array/mod.rs | 39 ++++++++++++++++--- .../core/src/iter/adapters/array_chunks.rs | 8 ++-- 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index 24e5d0fba21bf..2090756d7a3ec 100644 --- a/library/core/src/array/mod.rs +++ b/library/core/src/array/mod.rs @@ -878,13 +878,11 @@ where ControlFlow::Continue(elem) => elem, }; - // SAFETY: `guard.initialized` starts at 0, is increased by one in the - // loop and the loop is aborted once it reaches N (which is - // `array.len()`). + // SAFETY: `guard.initialized` starts at 0, which means push can be called + // at most N times, which this loop does. unsafe { - guard.array_mut.get_unchecked_mut(guard.initialized).write(item); + guard.push_unchecked(item); } - guard.initialized += 1; } None => { let alive = 0..guard.initialized; @@ -902,11 +900,42 @@ where Ok(Try::from_output(output)) } +/// Panic guard for incremental initialization of arrays. +/// +/// Disarm the guard with `mem::forget` once the array has been initialized. +/// +/// # Safety +/// +/// All write accesses to this structure are unsafe and must maintain a correct +/// count of `initialized` elements. +/// +/// To minimize indirection fields are still pub but callers should at least use +/// `push_unchecked` to signal that something unsafe is going on. pub(crate) struct Guard<'a, T, const N: usize> { + /// The array to be initialized. pub array_mut: &'a mut [MaybeUninit; N], + /// The number of items that have been initialized so far. pub initialized: usize, } +impl Guard<'_, T, N> { + /// Adds an item to the array and updates the initialized item counter. + /// + /// # Safety + /// + /// No more than N elements must be initialized. + #[inline] + pub unsafe fn push_unchecked(&mut self, item: T) { + // SAFETY: If `initialized` was correct before and the caller does not + // invoke this method more than N times then writes will be in-bounds + // and slots will not be initialized more than once. + unsafe { + self.array_mut.get_unchecked_mut(self.initialized).write(item); + self.initialized = self.initialized.unchecked_add(1); + } + } +} + impl Drop for Guard<'_, T, N> { fn drop(&mut self) { debug_assert!(self.initialized <= N); diff --git a/library/core/src/iter/adapters/array_chunks.rs b/library/core/src/iter/adapters/array_chunks.rs index 9bd135007b4a6..5e4211058aa6f 100644 --- a/library/core/src/iter/adapters/array_chunks.rs +++ b/library/core/src/iter/adapters/array_chunks.rs @@ -216,11 +216,13 @@ where while inner_len - i >= N { let mut chunk = MaybeUninit::uninit_array(); let mut guard = array::Guard { array_mut: &mut chunk, initialized: 0 }; - for j in 0..N { + while guard.initialized < N { // SAFETY: The method consumes the iterator and the loop condition ensures that // all accesses are in bounds and only happen once. - guard.array_mut[j].write(unsafe { self.iter.__iterator_get_unchecked(i + j) }); - guard.initialized = j + 1; + unsafe { + let idx = i + guard.initialized; + guard.push_unchecked(self.iter.__iterator_get_unchecked(idx)); + } } mem::forget(guard); // SAFETY: The loop above initialized all elements