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#[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    #[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                    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        self.next_back_remainder();
129
130        let mut acc = init;
131        let mut iter = ByRefSized(&mut self.iter).rev();
132
133        while let Ok(mut chunk) = iter.next_chunk() {
137            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    fn next_back_remainder(&mut self) {
155        if self.remainder.is_some() {
158            return;
159        }
160
161        let rem = self.iter.len() % N;
164
165        let mut remainder =
167            unsafe { self.iter.by_ref().rev().take(rem).next_chunk().unwrap_err_unchecked() };
169
170        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        while inner_len - i >= N {
234            let chunk = crate::array::from_fn(|local| {
235                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        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        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}