diff options
Diffstat (limited to '')
-rw-r--r-- | library/core/src/iter/adapters/array_chunks.rs | 75 |
1 files changed, 72 insertions, 3 deletions
diff --git a/library/core/src/iter/adapters/array_chunks.rs b/library/core/src/iter/adapters/array_chunks.rs index d4fb88610..5e4211058 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<B, F>(self, init: B, f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B, + { + <Self as SpecFold>::fold(self, init, f) + } } #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")] @@ -168,3 +176,64 @@ where self.iter.len() < N } } + +trait SpecFold: Iterator { + fn fold<B, F>(self, init: B, f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B; +} + +impl<I, const N: usize> SpecFold for ArrayChunks<I, N> +where + I: Iterator, +{ + #[inline] + default fn fold<B, F>(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<I, const N: usize> SpecFold for ArrayChunks<I, N> +where + I: Iterator + TrustedRandomAccessNoCoerce, +{ + #[inline] + fn fold<B, F>(mut self, init: B, mut f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B, + { + 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 }; + 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. + 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 + let chunk = unsafe { MaybeUninit::array_assume_init(chunk) }; + accum = f(accum, chunk); + i += N; + } + + // unlike try_fold this method does not need to take care of the remainder + // since `self` will be dropped + + accum + } +} |