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