summaryrefslogtreecommitdiffstats
path: root/library/core/src/iter/adapters/array_chunks.rs
diff options
context:
space:
mode:
Diffstat (limited to 'library/core/src/iter/adapters/array_chunks.rs')
-rw-r--r--library/core/src/iter/adapters/array_chunks.rs75
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
+ }
+}