core/iter/adapters/
array_chunks.rs

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 = self.iter.len() % N;
164
165        // Take the last `rem` elements out of `self.iter`.
166        let mut remainder =
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, 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), Some(n)) => m.checked_mul(n),
273            _ => None,
274        }
275    };
276}