1use crate::array;
2use crate::iter::adapters::SourceIter;
3use crate::iter::{
4 ByRefSized, FusedIterator, InPlaceIterable, TrustedFused, TrustedRandomAccessNoCoerce,
5};
6use crate::num::NonZero;
7use crate::ops::{ControlFlow, NeverShortCircuit, Try};
8
9#[derive(Debug, Clone)]
17#[must_use = "iterators are lazy and do nothing unless consumed"]
18#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
19pub struct ArrayChunks<I: Iterator, const N: usize> {
20 iter: I,
21 remainder: Option<array::IntoIter<I::Item, N>>,
22}
23
24impl<I, const N: usize> ArrayChunks<I, N>
25where
26 I: Iterator,
27{
28 #[track_caller]
29 pub(in crate::iter) fn new(iter: I) -> Self {
30 assert!(N != 0, "chunk size must be non-zero");
31 Self { iter, remainder: None }
32 }
33
34 #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
48 #[inline]
49 pub fn into_remainder(mut self) -> Option<array::IntoIter<I::Item, N>> {
50 if self.remainder.is_none() {
51 while let Some(_) = self.next() {}
52 }
53 self.remainder
54 }
55}
56
57#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
58impl<I, const N: usize> Iterator for ArrayChunks<I, N>
59where
60 I: Iterator,
61{
62 type Item = [I::Item; N];
63
64 #[inline]
65 fn next(&mut self) -> Option<Self::Item> {
66 self.try_for_each(ControlFlow::Break).break_value()
67 }
68
69 #[inline]
70 fn size_hint(&self) -> (usize, Option<usize>) {
71 let (lower, upper) = self.iter.size_hint();
72
73 (lower / N, upper.map(|n| n / N))
74 }
75
76 #[inline]
77 fn count(self) -> usize {
78 self.iter.count() / N
79 }
80
81 fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
82 where
83 Self: Sized,
84 F: FnMut(B, Self::Item) -> R,
85 R: Try<Output = B>,
86 {
87 let mut acc = init;
88 loop {
89 match self.iter.next_chunk() {
90 Ok(chunk) => acc = f(acc, chunk)?,
91 Err(remainder) => {
92 self.remainder.get_or_insert(remainder);
95
96 break try { acc };
97 }
98 }
99 }
100 }
101
102 fn fold<B, F>(self, init: B, f: F) -> B
103 where
104 Self: Sized,
105 F: FnMut(B, Self::Item) -> B,
106 {
107 <Self as SpecFold>::fold(self, init, f)
108 }
109}
110
111#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
112impl<I, const N: usize> DoubleEndedIterator for ArrayChunks<I, N>
113where
114 I: DoubleEndedIterator + ExactSizeIterator,
115{
116 #[inline]
117 fn next_back(&mut self) -> Option<Self::Item> {
118 self.try_rfold((), |(), x| ControlFlow::Break(x)).break_value()
119 }
120
121 fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
122 where
123 Self: Sized,
124 F: FnMut(B, Self::Item) -> R,
125 R: Try<Output = B>,
126 {
127 self.next_back_remainder();
129
130 let mut acc = init;
131 let mut iter = ByRefSized(&mut self.iter).rev();
132
133 while let Ok(mut chunk) = iter.next_chunk() {
137 chunk.reverse();
140 acc = f(acc, chunk)?
141 }
142
143 try { acc }
144 }
145
146 impl_fold_via_try_fold! { rfold -> try_rfold }
147}
148
149impl<I, const N: usize> ArrayChunks<I, N>
150where
151 I: DoubleEndedIterator + ExactSizeIterator,
152{
153 fn next_back_remainder(&mut self) {
155 if self.remainder.is_some() {
158 return;
159 }
160
161 let rem = self.iter.len() % N;
164
165 let mut remainder =
167 unsafe { self.iter.by_ref().rev().take(rem).next_chunk().unwrap_err_unchecked() };
169
170 remainder.as_mut_slice().reverse();
172 self.remainder = Some(remainder);
173 }
174}
175
176#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
177impl<I, const N: usize> FusedIterator for ArrayChunks<I, N> where I: FusedIterator {}
178
179#[unstable(issue = "none", feature = "trusted_fused")]
180unsafe impl<I, const N: usize> TrustedFused for ArrayChunks<I, N> where I: TrustedFused + Iterator {}
181
182#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
183impl<I, const N: usize> ExactSizeIterator for ArrayChunks<I, N>
184where
185 I: ExactSizeIterator,
186{
187 #[inline]
188 fn len(&self) -> usize {
189 self.iter.len() / N
190 }
191
192 #[inline]
193 fn is_empty(&self) -> bool {
194 self.iter.len() < N
195 }
196}
197
198trait SpecFold: Iterator {
199 fn fold<B, F>(self, init: B, f: F) -> B
200 where
201 Self: Sized,
202 F: FnMut(B, Self::Item) -> B;
203}
204
205impl<I, const N: usize> SpecFold for ArrayChunks<I, N>
206where
207 I: Iterator,
208{
209 #[inline]
210 default fn fold<B, F>(mut self, init: B, f: F) -> B
211 where
212 Self: Sized,
213 F: FnMut(B, Self::Item) -> B,
214 {
215 self.try_fold(init, NeverShortCircuit::wrap_mut_2(f)).0
216 }
217}
218
219impl<I, const N: usize> SpecFold for ArrayChunks<I, N>
220where
221 I: Iterator + TrustedRandomAccessNoCoerce,
222{
223 #[inline]
224 fn fold<B, F>(mut self, init: B, mut f: F) -> B
225 where
226 Self: Sized,
227 F: FnMut(B, Self::Item) -> B,
228 {
229 let mut accum = init;
230 let inner_len = self.iter.size();
231 let mut i = 0;
232 while inner_len - i >= N {
234 let chunk = crate::array::from_fn(|local| {
235 unsafe {
238 let idx = i + local;
239 self.iter.__iterator_get_unchecked(idx)
240 }
241 });
242 accum = f(accum, chunk);
243 i += N;
244 }
245
246 accum
250 }
251}
252
253#[unstable(issue = "none", feature = "inplace_iteration")]
254unsafe impl<I, const N: usize> SourceIter for ArrayChunks<I, N>
255where
256 I: SourceIter + Iterator,
257{
258 type Source = I::Source;
259
260 #[inline]
261 unsafe fn as_inner(&mut self) -> &mut I::Source {
262 unsafe { SourceIter::as_inner(&mut self.iter) }
264 }
265}
266
267#[unstable(issue = "none", feature = "inplace_iteration")]
268unsafe impl<I: InPlaceIterable + Iterator, const N: usize> InPlaceIterable for ArrayChunks<I, N> {
269 const EXPAND_BY: Option<NonZero<usize>> = I::EXPAND_BY;
270 const MERGE_BY: Option<NonZero<usize>> = const {
271 match (I::MERGE_BY, NonZero::new(N)) {
272 (Some(m), Some(n)) => m.checked_mul(n),
273 _ => None,
274 }
275 };
276}