array_chunks.rs - source
Description: Source of the Rust file `library/core/src/iter/adapters/array_chunks.rs`.

#
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) 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    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", 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", 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}