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

   1 use crate::array;
   2 use crate::iter::{ByRefSized, FusedIterator, Iterator, TrustedRandomAccessNoCoerce};
   3 use crate::mem::{self, MaybeUninit};
   4 use crate::ops::{ControlFlow, NeverShortCircuit, Try};
   5
   6 /// An iterator over `N` elements of the iterator at a time.
   7 ///
   8 /// The chunks do not overlap. If `N` does not divide the length of the
   9 /// iterator, then the last up to `N-1` elements will be omitted.
  10 ///
  11 /// This `struct` is created by the [`array_chunks`][Iterator::array_chunks]
  12 /// method on [`Iterator`]. See its documentation for more.
  13 #[derive(Debug, Clone)]
  14 #[must_use = "iterators are lazy and do nothing unless consumed"]
  15 #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
  16 pub struct ArrayChunks<I: Iterator, const N: usize> {
  17     iter: I,
  18     remainder: Option<array::IntoIter<I::Item, N>>,
  19 }
  20
  21 impl<I, const N: usize> ArrayChunks<I, N>
  22 where
  23     I: Iterator,
  24 {
  25     #[track_caller]
  26     pub(in crate::iter) fn new(iter: I) -> Self {
  27         assert!(N != 0, "chunk size must be non-zero");
  28         Self { iter, remainder: None }
  29     }
  30
  31     /// Returns an iterator over the remaining elements of the original iterator
  32     /// that are not going to be returned by this iterator. The returned
  33     /// iterator will yield at most `N-1` elements.
  34     #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
  35     #[inline]
  36     pub fn into_remainder(self) -> Option<array::IntoIter<I::Item, N>> {
  37         self.remainder
  38     }
  39 }
  40
  41 #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
  42 impl<I, const N: usize> Iterator for ArrayChunks<I, N>
  43 where
  44     I: Iterator,
  45 {
  46     type Item = [I::Item; N];
  47
  48     #[inline]
  49     fn next(&mut self) -> Option<Self::Item> {
  50         self.try_for_each(ControlFlow::Break).break_value()
  51     }
  52
  53     #[inline]
  54     fn size_hint(&self) -> (usize, Option<usize>) {
  55         let (lower, upper) = self.iter.size_hint();
  56
  57         (lower / N, upper.map(|n| n / N))
  58     }
  59
  60     #[inline]
  61     fn count(self) -> usize {
  62         self.iter.count() / N
  63     }
  64
  65     fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
  66     where
  67         Self: Sized,
  68         F: FnMut(B, Self::Item) -> R,
  69         R: Try<Output = B>,
  70     {
  71         let mut acc = init;
  72         loop {
  73             match self.iter.next_chunk() {
  74                 Ok(chunk) => acc = f(acc, chunk)?,
  75                 Err(remainder) => {
  76                     // Make sure to not override `self.remainder` with an empty array
  77                     // when `next` is called after `ArrayChunks` exhaustion.
  78                     self.remainder.get_or_insert(remainder);
  79
  80                     break try { acc };
  81                 }
  82             }
  83         }
  84     }
  85
  86     fn fold<B, F>(self, init: B, f: F) -> B
  87     where
  88         Self: Sized,
  89         F: FnMut(B, Self::Item) -> B,
  90     {
  91         <Self as SpecFold>::fold(self, init, f)
  92     }
  93 }
  94
  95 #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
  96 impl<I, const N: usize> DoubleEndedIterator for ArrayChunks<I, N>
  97 where
  98     I: DoubleEndedIterator + ExactSizeIterator,
  99 {
 100     #[inline]
 101     fn next_back(&mut self) -> Option<Self::Item> {
 102         self.try_rfold((), |(), x| ControlFlow::Break(x)).break_value()
 103     }
 104
 105     fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
 106     where
 107         Self: Sized,
 108         F: FnMut(B, Self::Item) -> R,
 109         R: Try<Output = B>,
 110     {
 111         // We are iterating from the back we need to first handle the remainder.
 112         self.next_back_remainder();
 113
 114         let mut acc = init;
 115         let mut iter = ByRefSized(&mut self.iter).rev();
 116
 117         // NB remainder is handled by `next_back_remainder`, so
 118         // `next_chunk` can't return `Err` with non-empty remainder
 119         // (assuming correct `I as ExactSizeIterator` impl).
 120         while let Ok(mut chunk) = iter.next_chunk() {
 121             // FIXME: do not do double reverse
 122             //        (we could instead add `next_chunk_back` for example)
 123             chunk.reverse();
 124             acc = f(acc, chunk)?
 125         }
 126
 127         try { acc }
 128     }
 129
 130     impl_fold_via_try_fold! { rfold -> try_rfold }
 131 }
 132
 133 impl<I, const N: usize> ArrayChunks<I, N>
 134 where
 135     I: DoubleEndedIterator + ExactSizeIterator,
 136 {
 137     /// Updates `self.remainder` such that `self.iter.len` is divisible by `N`.
 138     fn next_back_remainder(&mut self) {
 139         // Make sure to not override `self.remainder` with an empty array
 140         // when `next_back` is called after `ArrayChunks` exhaustion.
 141         if self.remainder.is_some() {
 142             return;
 143         }
 144
 145         // We use the `ExactSizeIterator` implementation of the underlying
 146         // iterator to know how many remaining elements there are.
 147         let rem = self.iter.len() % N;
 148
 149         // Take the last `rem` elements out of `self.iter`.
 150         let mut remainder =
 151             // SAFETY: `unwrap_err` always succeeds because x % N < N for all x.
 152             unsafe { self.iter.by_ref().rev().take(rem).next_chunk().unwrap_err_unchecked() };
 153
 154         // We used `.rev()` above, so we need to re-reverse the reminder
 155         remainder.as_mut_slice().reverse();
 156         self.remainder = Some(remainder);
 157     }
 158 }
 159
 160 #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
 161 impl<I, const N: usize> FusedIterator for ArrayChunks<I, N> where I: FusedIterator {}
 162
 163 #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "100450")]
 164 impl<I, const N: usize> ExactSizeIterator for ArrayChunks<I, N>
 165 where
 166     I: ExactSizeIterator,
 167 {
 168     #[inline]
 169     fn len(&self) -> usize {
 170         self.iter.len() / N
 171     }
 172
 173     #[inline]
 174     fn is_empty(&self) -> bool {
 175         self.iter.len() < N
 176     }
 177 }
 178
 179 trait SpecFold: Iterator {
 180     fn fold<B, F>(self, init: B, f: F) -> B
 181     where
 182         Self: Sized,
 183         F: FnMut(B, Self::Item) -> B;
 184 }
 185
 186 impl<I, const N: usize> SpecFold for ArrayChunks<I, N>
 187 where
 188     I: Iterator,
 189 {
 190     #[inline]
 191     default fn fold<B, F>(mut self, init: B, f: F) -> B
 192     where
 193         Self: Sized,
 194         F: FnMut(B, Self::Item) -> B,
 195     {
 196         self.try_fold(init, NeverShortCircuit::wrap_mut_2(f)).0
 197     }
 198 }
 199
 200 impl<I, const N: usize> SpecFold for ArrayChunks<I, N>
 201 where
 202     I: Iterator + TrustedRandomAccessNoCoerce,
 203 {
 204     #[inline]
 205     fn fold<B, F>(mut self, init: B, mut f: F) -> B
 206     where
 207         Self: Sized,
 208         F: FnMut(B, Self::Item) -> B,
 209     {
 210         let mut accum = init;
 211         let inner_len = self.iter.size();
 212         let mut i = 0;
 213         // Use a while loop because (0..len).step_by(N) doesn't optimize well.
 214         while inner_len - i >= N {
 215             let mut chunk = MaybeUninit::uninit_array();
 216             let mut guard = array::Guard { array_mut: &mut chunk, initialized: 0 };
 217             while guard.initialized < N {
 218                 // SAFETY: The method consumes the iterator and the loop condition ensures that
 219                 // all accesses are in bounds and only happen once.
 220                 unsafe {
 221                     let idx = i + guard.initialized;
 222                     guard.push_unchecked(self.iter.__iterator_get_unchecked(idx));
 223                 }
 224             }
 225             mem::forget(guard);
 226             // SAFETY: The loop above initialized all elements
 227             let chunk = unsafe { MaybeUninit::array_assume_init(chunk) };
 228             accum = f(accum, chunk);
 229             i += N;
 230         }
 231
 232         // unlike try_fold this method does not need to take care of the remainder
 233         // since `self` will be dropped
 234
 235         accum
 236     }
 237 }