library/core/src/iter/adapters/array_chunks.rs

   1 use crate::array;
   2 use crate::iter::{ByRefSized, FusedIterator, Iterator};
   3 use crate::ops::{ControlFlow, Try};
   4
   5 /// An iterator over `N` elements of the iterator at a time.
   6 ///
   7 /// The chunks do not overlap. If `N` does not divide the length of the
   8 /// iterator, then the last up to `N-1` elements will be omitted.
   9 ///
  10 /// This `struct` is created by the [`array_chunks`][Iterator::array_chunks]
  11 /// method on [`Iterator`]. See its documentation for more.
  12 #[derive(Debug, Clone)]
  13 #[must_use = "iterators are lazy and do nothing unless consumed"]
  14 #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
  15 pub struct ArrayChunks<I: Iterator, const N: usize> {
  16     iter: I,
  17     remainder: Option<array::IntoIter<I::Item, N>>,
  18 }
  19
  20 impl<I, const N: usize> ArrayChunks<I, N>
  21 where
  22     I: Iterator,
  23 {
  24     #[track_caller]
  25     pub(in crate::iter) fn new(iter: I) -> Self {
  26         assert!(N != 0, "chunk size must be non-zero");
  27         Self { iter, remainder: None }
  28     }
  29
  30     /// Returns an iterator over the remaining elements of the original iterator
  31     /// that are not going to be returned by this iterator. The returned
  32     /// iterator will yield at most `N-1` elements.
  33     #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
  34     #[inline]
  35     pub fn into_remainder(self) -> Option<array::IntoIter<I::Item, N>> {
  36         self.remainder
  37     }
  38 }
  39
  40 #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
  41 impl<I, const N: usize> Iterator for ArrayChunks<I, N>
  42 where
  43     I: Iterator,
  44 {
  45     type Item = [I::Item; N];
  46
  47     #[inline]
  48     fn next(&mut self) -> Option<Self::Item> {
  49         self.try_for_each(ControlFlow::Break).break_value()
  50     }
  51
  52     #[inline]
  53     fn size_hint(&self) -> (usize, Option<usize>) {
  54         let (lower, upper) = self.iter.size_hint();
  55
  56         (lower / N, upper.map(|n| n / N))
  57     }
  58
  59     #[inline]
  60     fn count(self) -> usize {
  61         self.iter.count() / N
  62     }
  63
  64     fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
  65     where
  66         Self: Sized,
  67         F: FnMut(B, Self::Item) -> R,
  68         R: Try<Output = B>,
  69     {
  70         let mut acc = init;
  71         loop {
  72             match self.iter.next_chunk() {
  73                 Ok(chunk) => acc = f(acc, chunk)?,
  74                 Err(remainder) => {
  75                     // Make sure to not override `self.remainder` with an empty array
  76                     // when `next` is called after `ArrayChunks` exhaustion.
  77                     self.remainder.get_or_insert(remainder);
  78
  79                     break try { acc };
  80                 }
  81             }
  82         }
  83     }
  84
  85     impl_fold_via_try_fold! { fold -> try_fold }
  86 }
  87
  88 #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
  89 impl<I, const N: usize> DoubleEndedIterator for ArrayChunks<I, N>
  90 where
  91     I: DoubleEndedIterator + ExactSizeIterator,
  92 {
  93     #[inline]
  94     fn next_back(&mut self) -> Option<Self::Item> {
  95         self.try_rfold((), |(), x| ControlFlow::Break(x)).break_value()
  96     }
  97
  98     fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
  99     where
 100         Self: Sized,
 101         F: FnMut(B, Self::Item) -> R,
 102         R: Try<Output = B>,
 103     {
 104         // We are iterating from the back we need to first handle the remainder.
 105         self.next_back_remainder();
 106
 107         let mut acc = init;
 108         let mut iter = ByRefSized(&mut self.iter).rev();
 109
 110         // NB remainder is handled by `next_back_remainder`, so
 111         // `next_chunk` can't return `Err` with non-empty remainder
 112         // (assuming correct `I as ExactSizeIterator` impl).
 113         while let Ok(mut chunk) = iter.next_chunk() {
 114             // FIXME: do not do double reverse
 115             //        (we could instead add `next_chunk_back` for example)
 116             chunk.reverse();
 117             acc = f(acc, chunk)?
 118         }
 119
 120         try { acc }
 121     }
 122
 123     impl_fold_via_try_fold! { rfold -> try_rfold }
 124 }
 125
 126 impl<I, const N: usize> ArrayChunks<I, N>
 127 where
 128     I: DoubleEndedIterator + ExactSizeIterator,
 129 {
 130     /// Updates `self.remainder` such that `self.iter.len` is divisible by `N`.
 131     fn next_back_remainder(&mut self) {
 132         // Make sure to not override `self.remainder` with an empty array
 133         // when `next_back` is called after `ArrayChunks` exhaustion.
 134         if self.remainder.is_some() {
 135             return;
 136         }
 137
 138         // We use the `ExactSizeIterator` implementation of the underlying
 139         // iterator to know how many remaining elements there are.
 140         let rem = self.iter.len() % N;
 141
 142         // Take the last `rem` elements out of `self.iter`.
 143         let mut remainder =
 144             // SAFETY: `unwrap_err` always succeeds because x % N < N for all x.
 145             unsafe { self.iter.by_ref().rev().take(rem).next_chunk().unwrap_err_unchecked() };
 146
 147         // We used `.rev()` above, so we need to re-reverse the reminder
 148         remainder.as_mut_slice().reverse();
 149         self.remainder = Some(remainder);
 150     }
 151 }
 152
 153 #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
 154 impl<I, const N: usize> FusedIterator for ArrayChunks<I, N> where I: FusedIterator {}
 155
 156 #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
 157 impl<I, const N: usize> ExactSizeIterator for ArrayChunks<I, N>
 158 where
 159     I: ExactSizeIterator,
 160 {
 161     #[inline]
 162     fn len(&self) -> usize {
 163         self.iter.len() / N
 164     }
 165
 166     #[inline]
 167     fn is_empty(&self) -> bool {
 168         self.iter.len() < N
 169     }
 170 }