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/// 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 ///
38 /// # Example
39 /// ```
40 /// # // Also serves as a regression test for https://github.com/rust-lang/rust/issues/123333
41 /// # #![feature(iter_array_chunks)]
42 /// let x = [1,2,3,4,5].into_iter().array_chunks::<2>();
43 /// let mut rem = x.into_remainder().unwrap();
44 /// assert_eq!(rem.next(), Some(5));
45 /// assert_eq!(rem.next(), None);
46 /// ```
47 #[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 // Make sure to not override `self.remainder` with an empty array
93 // when `next` is called after `ArrayChunks` exhaustion.
94 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 // We are iterating from the back we need to first handle the remainder.
128 self.next_back_remainder();
129
130 let mut acc = init;
131 let mut iter = ByRefSized(&mut self.iter).rev();
132
133 // NB remainder is handled by `next_back_remainder`, so
134 // `next_chunk` can't return `Err` with non-empty remainder
135 // (assuming correct `I as ExactSizeIterator` impl).
136 while let Ok(mut chunk) = iter.next_chunk() {
137 // FIXME: do not do double reverse
138 // (we could instead add `next_chunk_back` for example)
139 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 /// Updates `self.remainder` such that `self.iter.len` is divisible by `N`.
154 fn next_back_remainder(&mut self) {
155 // Make sure to not override `self.remainder` with an empty array
156 // when `next_back` is called after `ArrayChunks` exhaustion.
157 if self.remainder.is_some() {
158 return;
159 }
160
161 // We use the `ExactSizeIterator` implementation of the underlying
162 // iterator to know how many remaining elements there are.
163 let rem: usize = self.iter.len() % N;
164
165 // Take the last `rem` elements out of `self.iter`.
166 let mut remainder: IntoIter<::Item, N> =
167 // SAFETY: `unwrap_err` always succeeds because x % N < N for all x.
168 unsafe { self.iter.by_ref().rev().take(rem).next_chunk().unwrap_err_unchecked() };
169
170 // We used `.rev()` above, so we need to re-reverse the reminder
171 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, f: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 // Use a while loop because (0..len).step_by(N) doesn't optimize well.
233 while inner_len - i >= N {
234 let chunk = crate::array::from_fn(|local| {
235 // SAFETY: The method consumes the iterator and the loop condition ensures that
236 // all accesses are in bounds and only happen once.
237 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 // unlike try_fold this method does not need to take care of the remainder
247 // since `self` will be dropped
248
249 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 // SAFETY: unsafe function forwarding to unsafe function with the same requirements
263 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: NonZero), Some(n: NonZero)) => m.checked_mul(n),
273 _ => None,
274 }
275 };
276}
277