core/iter/adapters/
array_chunks.rs

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

core/iter/adapters/
array_chunks.rs
