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    fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
120    where
121        Self: Sized,
122        F: FnMut(B, Self::Item) -> R,
123        R: Try<Output = B>,
124    {
125        // We are iterating from the back we need to first handle the remainder.
126        self.next_back_remainder();
127
128        let mut acc = init;
129
130        // NB remainder is handled by `next_back_remainder`, so
131        // `next_chunk_back` can't return `Err` with non-empty remainder
132        // (assuming correct `I as ExactSizeIterator` impl).
133        while let Ok(chunk) = self.iter.next_chunk_back() {
134            acc = f(acc, chunk)?
135        }
136
137        try { acc }
138    }
139
140    impl_fold_via_try_fold! { rfold -> try_rfold }
141}
142
143impl<I, const N: usize> ArrayChunks<I, N>
144where
145    I: DoubleEndedIterator + ExactSizeIterator,
146{
147    /// Updates `self.remainder` such that `self.iter.len` is divisible by `N`.
148    fn next_back_remainder(&mut self) {
149        // Make sure to not override `self.remainder` with an empty array
150        // when `next_back` is called after `ArrayChunks` exhaustion.
151        if self.remainder.is_some() {
152            return;
153        }
154
155        // We use the `ExactSizeIterator` implementation of the underlying
156        // iterator to know how many remaining elements there are.
157        let rem = self.iter.len() % N;
158
159        // Take the last `rem` elements out of `self.iter`.
160        let mut remainder =
161            // SAFETY: `unwrap_err` always succeeds because x % N < N for all x.
162            unsafe { self.iter.by_ref().rev().take(rem).next_chunk().unwrap_err_unchecked() };
163
164        // We used `.rev()` above, so we need to re-reverse the reminder
165        remainder.as_mut_slice().reverse();
166        self.remainder = Some(remainder);
167    }
168}
169
170#[unstable(feature = "iter_array_chunks", issue = "100450")]
171impl<I, const N: usize> FusedIterator for ArrayChunks<I, N> where I: FusedIterator {}
172
173#[unstable(issue = "none", feature = "trusted_fused")]
174unsafe impl<I, const N: usize> TrustedFused for ArrayChunks<I, N> where I: TrustedFused + Iterator {}
175
176#[unstable(feature = "iter_array_chunks", issue = "100450")]
177impl<I, const N: usize> ExactSizeIterator for ArrayChunks<I, N>
178where
179    I: ExactSizeIterator,
180{
181    #[inline]
182    fn len(&self) -> usize {
183        self.iter.len() / N
184    }
185
186    #[inline]
187    fn is_empty(&self) -> bool {
188        self.iter.len() < N
189    }
190}
191
192trait SpecFold: Iterator {
193    fn fold<B, F>(self, init: B, f: F) -> B
194    where
195        Self: Sized,
196        F: FnMut(B, Self::Item) -> B;
197}
198
199impl<I, const N: usize> SpecFold for ArrayChunks<I, N>
200where
201    I: Iterator,
202{
203    #[inline]
204    default fn fold<B, F>(mut self, init: B, f: F) -> B
205    where
206        Self: Sized,
207        F: FnMut(B, Self::Item) -> B,
208    {
209        self.try_fold(init, NeverShortCircuit::wrap_mut_2(f)).0
210    }
211}
212
213impl<I, const N: usize> SpecFold for ArrayChunks<I, N>
214where
215    I: Iterator + TrustedRandomAccessNoCoerce,
216{
217    #[inline]
218    fn fold<B, F>(mut self, init: B, mut f: F) -> B
219    where
220        Self: Sized,
221        F: FnMut(B, Self::Item) -> B,
222    {
223        let mut accum = init;
224        let inner_len = self.iter.size();
225        let mut i = 0;
226        // Use a while loop because (0..len).step_by(N) doesn't optimize well.
227        while inner_len - i >= N {
228            let chunk = crate::array::from_fn(|local| {
229                // SAFETY: The method consumes the iterator and the loop condition ensures that
230                // all accesses are in bounds and only happen once.
231                unsafe {
232                    let idx = i + local;
233                    self.iter.__iterator_get_unchecked(idx)
234                }
235            });
236            accum = f(accum, chunk);
237            i += N;
238        }
239
240        // unlike try_fold this method does not need to take care of the remainder
241        // since `self` will be dropped
242
243        accum
244    }
245}
246
247#[unstable(issue = "none", feature = "inplace_iteration")]
248unsafe impl<I, const N: usize> SourceIter for ArrayChunks<I, N>
249where
250    I: SourceIter + Iterator,
251{
252    type Source = I::Source;
253
254    #[inline]
255    unsafe fn as_inner(&mut self) -> &mut I::Source {
256        // SAFETY: unsafe function forwarding to unsafe function with the same requirements
257        unsafe { SourceIter::as_inner(&mut self.iter) }
258    }
259}
260
261#[unstable(issue = "none", feature = "inplace_iteration")]
262unsafe impl<I: InPlaceIterable + Iterator, const N: usize> InPlaceIterable for ArrayChunks<I, N> {
263    const EXPAND_BY: Option<NonZero<usize>> = I::EXPAND_BY;
264    const MERGE_BY: Option<NonZero<usize>> = const {
265        match (I::MERGE_BY, NonZero::new(N)) {
266            (Some(m), Some(n)) => m.checked_mul(n),
267            _ => None,
268        }
269    };
270}
core/iter/adapters/array_chunks.rs

core/iter/adapters/
array_chunks.rs
