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

core/iter/adapters/
array_chunks.rs
