1 | use crate::array; |
2 | use crate::iter::adapters::SourceIter; |
3 | use crate::iter::{ |
4 | ByRefSized, FusedIterator, InPlaceIterable, TrustedFused, TrustedRandomAccessNoCoerce, |
5 | }; |
6 | use crate::num::NonZero; |
7 | use 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" )] |
19 | pub struct ArrayChunks<I: Iterator, const N: usize> { |
20 | iter: I, |
21 | remainder: Option<array::IntoIter<I::Item, N>>, |
22 | } |
23 | |
24 | impl<I, const N: usize> ArrayChunks<I, N> |
25 | where |
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" )] |
58 | impl<I, const N: usize> Iterator for ArrayChunks<I, N> |
59 | where |
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" )] |
112 | impl<I, const N: usize> DoubleEndedIterator for ArrayChunks<I, N> |
113 | where |
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 | |
149 | impl<I, const N: usize> ArrayChunks<I, N> |
150 | where |
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" )] |
177 | impl<I, const N: usize> FusedIterator for ArrayChunks<I, N> where I: FusedIterator {} |
178 | |
179 | #[unstable (issue = "none" , feature = "trusted_fused" )] |
180 | unsafe 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" )] |
183 | impl<I, const N: usize> ExactSizeIterator for ArrayChunks<I, N> |
184 | where |
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 | |
198 | trait 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 | |
205 | impl<I, const N: usize> SpecFold for ArrayChunks<I, N> |
206 | where |
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 | |
219 | impl<I, const N: usize> SpecFold for ArrayChunks<I, N> |
220 | where |
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" )] |
254 | unsafe impl<I, const N: usize> SourceIter for ArrayChunks<I, N> |
255 | where |
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" )] |
268 | unsafe 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 | |