1use crate::array;
2use crate::iter::adapters::SourceIter;
3use crate::iter::{
4 ByRefSized, FusedIterator, InPlaceIterable, TrustedFused, TrustedRandomAccessNoCoerce,
5};
6use crate::num::NonZeroUsize;
7use crate::ops::{ControlFlow, NeverShortCircuit, Try};
8
9/// An iterator over `N` elements of the iterator at a time.
10///
11/// The chunks do not overlap. If `N` does not divide the length of the
12/// iterator, then the last up to `N-1` elements will be omitted.
13///
14/// This `struct` is created by the [`array_chunks`][Iterator::array_chunks]
15/// method on [`Iterator`]. See its documentation for more.
16#[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 /// Returns an iterator over the remaining elements of the original iterator
35 /// that are not going to be returned by this iterator. The returned
36 /// iterator will yield at most `N-1` elements.
37 #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
38 #[inline]
39 pub fn into_remainder(self) -> Option<array::IntoIter<I::Item, N>> {
40 self.remainder
41 }
42}
43
44#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
45impl<I, const N: usize> Iterator for ArrayChunks<I, N>
46where
47 I: Iterator,
48{
49 type Item = [I::Item; N];
50
51 #[inline]
52 fn next(&mut self) -> Option<Self::Item> {
53 self.try_for_each(ControlFlow::Break).break_value()
54 }
55
56 #[inline]
57 fn size_hint(&self) -> (usize, Option<usize>) {
58 let (lower, upper) = self.iter.size_hint();
59
60 (lower / N, upper.map(|n| n / N))
61 }
62
63 #[inline]
64 fn count(self) -> usize {
65 self.iter.count() / N
66 }
67
68 fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
69 where
70 Self: Sized,
71 F: FnMut(B, Self::Item) -> R,
72 R: Try<Output = B>,
73 {
74 let mut acc = init;
75 loop {
76 match self.iter.next_chunk() {
77 Ok(chunk) => acc = f(acc, chunk)?,
78 Err(remainder) => {
79 // Make sure to not override `self.remainder` with an empty array
80 // when `next` is called after `ArrayChunks` exhaustion.
81 self.remainder.get_or_insert(remainder);
82
83 break try { acc };
84 }
85 }
86 }
87 }
88
89 fn fold<B, F>(self, init: B, f: F) -> B
90 where
91 Self: Sized,
92 F: FnMut(B, Self::Item) -> B,
93 {
94 <Self as SpecFold>::fold(self, init, f)
95 }
96}
97
98#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
99impl<I, const N: usize> DoubleEndedIterator for ArrayChunks<I, N>
100where
101 I: DoubleEndedIterator + ExactSizeIterator,
102{
103 #[inline]
104 fn next_back(&mut self) -> Option<Self::Item> {
105 self.try_rfold((), |(), x| ControlFlow::Break(x)).break_value()
106 }
107
108 fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
109 where
110 Self: Sized,
111 F: FnMut(B, Self::Item) -> R,
112 R: Try<Output = B>,
113 {
114 // We are iterating from the back we need to first handle the remainder.
115 self.next_back_remainder();
116
117 let mut acc = init;
118 let mut iter = ByRefSized(&mut self.iter).rev();
119
120 // NB remainder is handled by `next_back_remainder`, so
121 // `next_chunk` can't return `Err` with non-empty remainder
122 // (assuming correct `I as ExactSizeIterator` impl).
123 while let Ok(mut chunk) = iter.next_chunk() {
124 // FIXME: do not do double reverse
125 // (we could instead add `next_chunk_back` for example)
126 chunk.reverse();
127 acc = f(acc, chunk)?
128 }
129
130 try { acc }
131 }
132
133 impl_fold_via_try_fold! { rfold -> try_rfold }
134}
135
136impl<I, const N: usize> ArrayChunks<I, N>
137where
138 I: DoubleEndedIterator + ExactSizeIterator,
139{
140 /// Updates `self.remainder` such that `self.iter.len` is divisible by `N`.
141 fn next_back_remainder(&mut self) {
142 // Make sure to not override `self.remainder` with an empty array
143 // when `next_back` is called after `ArrayChunks` exhaustion.
144 if self.remainder.is_some() {
145 return;
146 }
147
148 // We use the `ExactSizeIterator` implementation of the underlying
149 // iterator to know how many remaining elements there are.
150 let rem: usize = self.iter.len() % N;
151
152 // Take the last `rem` elements out of `self.iter`.
153 let mut remainder: IntoIter<::Item, N> =
154 // SAFETY: `unwrap_err` always succeeds because x % N < N for all x.
155 unsafe { self.iter.by_ref().rev().take(rem).next_chunk().unwrap_err_unchecked() };
156
157 // We used `.rev()` above, so we need to re-reverse the reminder
158 remainder.as_mut_slice().reverse();
159 self.remainder = Some(remainder);
160 }
161}
162
163#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
164impl<I, const N: usize> FusedIterator for ArrayChunks<I, N> where I: FusedIterator {}
165
166#[unstable(issue = "none", feature = "trusted_fused")]
167unsafe impl<I, const N: usize> TrustedFused for ArrayChunks<I, N> where I: TrustedFused + Iterator {}
168
169#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
170impl<I, const N: usize> ExactSizeIterator for ArrayChunks<I, N>
171where
172 I: ExactSizeIterator,
173{
174 #[inline]
175 fn len(&self) -> usize {
176 self.iter.len() / N
177 }
178
179 #[inline]
180 fn is_empty(&self) -> bool {
181 self.iter.len() < N
182 }
183}
184
185trait SpecFold: Iterator {
186 fn fold<B, F>(self, init: B, f: F) -> B
187 where
188 Self: Sized,
189 F: FnMut(B, Self::Item) -> B;
190}
191
192impl<I, const N: usize> SpecFold for ArrayChunks<I, N>
193where
194 I: Iterator,
195{
196 #[inline]
197 default fn fold<B, F>(mut self, init: B, f: F) -> B
198 where
199 Self: Sized,
200 F: FnMut(B, Self::Item) -> B,
201 {
202 self.try_fold(init, f:NeverShortCircuit::wrap_mut_2(f)).0
203 }
204}
205
206impl<I, const N: usize> SpecFold for ArrayChunks<I, N>
207where
208 I: Iterator + TrustedRandomAccessNoCoerce,
209{
210 #[inline]
211 fn fold<B, F>(mut self, init: B, mut f: F) -> B
212 where
213 Self: Sized,
214 F: FnMut(B, Self::Item) -> B,
215 {
216 let mut accum = init;
217 let inner_len = self.iter.size();
218 let mut i = 0;
219 // Use a while loop because (0..len).step_by(N) doesn't optimize well.
220 while inner_len - i >= N {
221 let chunk = crate::array::from_fn(|local| {
222 // SAFETY: The method consumes the iterator and the loop condition ensures that
223 // all accesses are in bounds and only happen once.
224 unsafe {
225 let idx = i + local;
226 self.iter.__iterator_get_unchecked(idx)
227 }
228 });
229 accum = f(accum, chunk);
230 i += N;
231 }
232
233 // unlike try_fold this method does not need to take care of the remainder
234 // since `self` will be dropped
235
236 accum
237 }
238}
239
240#[unstable(issue = "none", feature = "inplace_iteration")]
241unsafe impl<I, const N: usize> SourceIter for ArrayChunks<I, N>
242where
243 I: SourceIter + Iterator,
244{
245 type Source = I::Source;
246
247 #[inline]
248 unsafe fn as_inner(&mut self) -> &mut I::Source {
249 // SAFETY: unsafe function forwarding to unsafe function with the same requirements
250 unsafe { SourceIter::as_inner(&mut self.iter) }
251 }
252}
253
254#[unstable(issue = "none", feature = "inplace_iteration")]
255unsafe impl<I: InPlaceIterable + Iterator, const N: usize> InPlaceIterable for ArrayChunks<I, N> {
256 const EXPAND_BY: Option<NonZeroUsize> = I::EXPAND_BY;
257 const MERGE_BY: Option<NonZeroUsize> = const {
258 match (I::MERGE_BY, NonZeroUsize::new(N)) {
259 (Some(m: NonZero), Some(n: NonZero)) => m.checked_mul(n),
260 _ => None,
261 }
262 };
263}
264