vendor/rustc-rayon/src/iter/chunks.rs

   1 use std::cmp::min;
   2
   3 use super::plumbing::*;
   4 use super::*;
   5 use crate::math::div_round_up;
   6
   7 /// `Chunks` is an iterator that groups elements of an underlying iterator.
   8 ///
   9 /// This struct is created by the [`chunks()`] method on [`IndexedParallelIterator`]
  10 ///
  11 /// [`chunks()`]: trait.IndexedParallelIterator.html#method.chunks
  12 /// [`IndexedParallelIterator`]: trait.IndexedParallelIterator.html
  13 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
  14 #[derive(Debug, Clone)]
  15 pub struct Chunks<I>
  16 where
  17     I: IndexedParallelIterator,
  18 {
  19     size: usize,
  20     i: I,
  21 }
  22
  23 impl<I> Chunks<I>
  24 where
  25     I: IndexedParallelIterator,
  26 {
  27     /// Creates a new `Chunks` iterator
  28     pub(super) fn new(i: I, size: usize) -> Self {
  29         Chunks { i, size }
  30     }
  31 }
  32
  33 impl<I> ParallelIterator for Chunks<I>
  34 where
  35     I: IndexedParallelIterator,
  36 {
  37     type Item = Vec<I::Item>;
  38
  39     fn drive_unindexed<C>(self, consumer: C) -> C::Result
  40     where
  41         C: Consumer<Vec<I::Item>>,
  42     {
  43         bridge(self, consumer)
  44     }
  45
  46     fn opt_len(&self) -> Option<usize> {
  47         Some(self.len())
  48     }
  49 }
  50
  51 impl<I> IndexedParallelIterator for Chunks<I>
  52 where
  53     I: IndexedParallelIterator,
  54 {
  55     fn drive<C>(self, consumer: C) -> C::Result
  56     where
  57         C: Consumer<Self::Item>,
  58     {
  59         bridge(self, consumer)
  60     }
  61
  62     fn len(&self) -> usize {
  63         div_round_up(self.i.len(), self.size)
  64     }
  65
  66     fn with_producer<CB>(self, callback: CB) -> CB::Output
  67     where
  68         CB: ProducerCallback<Self::Item>,
  69     {
  70         let len = self.i.len();
  71         return self.i.with_producer(Callback {
  72             size: self.size,
  73             len,
  74             callback,
  75         });
  76
  77         struct Callback<CB> {
  78             size: usize,
  79             len: usize,
  80             callback: CB,
  81         }
  82
  83         impl<T, CB> ProducerCallback<T> for Callback<CB>
  84         where
  85             CB: ProducerCallback<Vec<T>>,
  86         {
  87             type Output = CB::Output;
  88
  89             fn callback<P>(self, base: P) -> CB::Output
  90             where
  91                 P: Producer<Item = T>,
  92             {
  93                 self.callback.callback(ChunkProducer {
  94                     chunk_size: self.size,
  95                     len: self.len,
  96                     base,
  97                 })
  98             }
  99         }
 100     }
 101 }
 102
 103 struct ChunkProducer<P>
 104 where
 105     P: Producer,
 106 {
 107     chunk_size: usize,
 108     len: usize,
 109     base: P,
 110 }
 111
 112 impl<P> Producer for ChunkProducer<P>
 113 where
 114     P: Producer,
 115 {
 116     type Item = Vec<P::Item>;
 117     type IntoIter = ChunkSeq<P>;
 118
 119     fn into_iter(self) -> Self::IntoIter {
 120         ChunkSeq {
 121             chunk_size: self.chunk_size,
 122             len: self.len,
 123             inner: if self.len > 0 { Some(self.base) } else { None },
 124         }
 125     }
 126
 127     fn split_at(self, index: usize) -> (Self, Self) {
 128         let elem_index = min(index * self.chunk_size, self.len);
 129         let (left, right) = self.base.split_at(elem_index);
 130         (
 131             ChunkProducer {
 132                 chunk_size: self.chunk_size,
 133                 len: elem_index,
 134                 base: left,
 135             },
 136             ChunkProducer {
 137                 chunk_size: self.chunk_size,
 138                 len: self.len - elem_index,
 139                 base: right,
 140             },
 141         )
 142     }
 143
 144     fn min_len(&self) -> usize {
 145         div_round_up(self.base.min_len(), self.chunk_size)
 146     }
 147
 148     fn max_len(&self) -> usize {
 149         self.base.max_len() / self.chunk_size
 150     }
 151 }
 152
 153 struct ChunkSeq<P> {
 154     chunk_size: usize,
 155     len: usize,
 156     inner: Option<P>,
 157 }
 158
 159 impl<P> Iterator for ChunkSeq<P>
 160 where
 161     P: Producer,
 162 {
 163     type Item = Vec<P::Item>;
 164
 165     fn next(&mut self) -> Option<Self::Item> {
 166         let producer = self.inner.take()?;
 167         if self.len > self.chunk_size {
 168             let (left, right) = producer.split_at(self.chunk_size);
 169             self.inner = Some(right);
 170             self.len -= self.chunk_size;
 171             Some(left.into_iter().collect())
 172         } else {
 173             debug_assert!(self.len > 0);
 174             self.len = 0;
 175             Some(producer.into_iter().collect())
 176         }
 177     }
 178
 179     fn size_hint(&self) -> (usize, Option<usize>) {
 180         let len = self.len();
 181         (len, Some(len))
 182     }
 183 }
 184
 185 impl<P> ExactSizeIterator for ChunkSeq<P>
 186 where
 187     P: Producer,
 188 {
 189     #[inline]
 190     fn len(&self) -> usize {
 191         div_round_up(self.len, self.chunk_size)
 192     }
 193 }
 194
 195 impl<P> DoubleEndedIterator for ChunkSeq<P>
 196 where
 197     P: Producer,
 198 {
 199     fn next_back(&mut self) -> Option<Self::Item> {
 200         let producer = self.inner.take()?;
 201         if self.len > self.chunk_size {
 202             let mut size = self.len % self.chunk_size;
 203             if size == 0 {
 204                 size = self.chunk_size;
 205             }
 206             let (left, right) = producer.split_at(self.len - size);
 207             self.inner = Some(left);
 208             self.len -= size;
 209             Some(right.into_iter().collect())
 210         } else {
 211             debug_assert!(self.len > 0);
 212             self.len = 0;
 213             Some(producer.into_iter().collect())
 214         }
 215     }
 216 }