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

   1 use super::plumbing::*;
   2 use super::*;
   3 use std::cmp;
   4 use std::iter;
   5 use rayon_core::join;
   6
   7 /// `Chain` is an iterator that joins `b` after `a` in one continuous iterator.
   8 /// This struct is created by the [`chain()`] method on [`ParallelIterator`]
   9 ///
  10 /// [`chain()`]: trait.ParallelIterator.html#method.chain
  11 /// [`ParallelIterator`]: trait.ParallelIterator.html
  12 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
  13 #[derive(Debug, Clone)]
  14 pub struct Chain<A, B>
  15     where A: ParallelIterator,
  16           B: ParallelIterator<Item = A::Item>
  17 {
  18     a: A,
  19     b: B,
  20 }
  21
  22 /// Create a new `Chain` iterator.
  23 ///
  24 /// NB: a free fn because it is NOT part of the end-user API.
  25 pub fn new<A, B>(a: A, b: B) -> Chain<A, B>
  26     where A: ParallelIterator,
  27           B: ParallelIterator<Item = A::Item>
  28 {
  29     Chain { a: a, b: b }
  30 }
  31
  32 impl<A, B> ParallelIterator for Chain<A, B>
  33     where A: ParallelIterator,
  34           B: ParallelIterator<Item = A::Item>
  35 {
  36     type Item = A::Item;
  37
  38     fn drive_unindexed<C>(self, consumer: C) -> C::Result
  39         where C: UnindexedConsumer<Self::Item>
  40     {
  41         let Chain { a, b } = self;
  42
  43         // If we returned a value from our own `opt_len`, then the collect consumer in particular
  44         // will balk at being treated like an actual `UnindexedConsumer`.  But when we do know the
  45         // length, we can use `Consumer::split_at` instead, and this is still harmless for other
  46         // truly-unindexed consumers too.
  47         let (left, right, reducer) = if let Some(len) = a.opt_len() {
  48             consumer.split_at(len)
  49         } else {
  50             let reducer = consumer.to_reducer();
  51             (consumer.split_off_left(), consumer, reducer)
  52         };
  53
  54         let (a, b) = join(|| a.drive_unindexed(left), || b.drive_unindexed(right));
  55         reducer.reduce(a, b)
  56     }
  57
  58     fn opt_len(&self) -> Option<usize> {
  59         match (self.a.opt_len(), self.b.opt_len()) {
  60             (Some(a_len), Some(b_len)) => a_len.checked_add(b_len),
  61             _ => None,
  62         }
  63     }
  64 }
  65
  66 impl<A, B> IndexedParallelIterator for Chain<A, B>
  67     where A: IndexedParallelIterator,
  68           B: IndexedParallelIterator<Item = A::Item>
  69 {
  70     fn drive<C>(self, consumer: C) -> C::Result
  71         where C: Consumer<Self::Item>
  72     {
  73         let Chain { a, b } = self;
  74         let (left, right, reducer) = consumer.split_at(a.len());
  75         let (a, b) = join(|| a.drive(left), || b.drive(right));
  76         reducer.reduce(a, b)
  77     }
  78
  79     fn len(&self) -> usize {
  80         self.a
  81             .len()
  82             .checked_add(self.b.len())
  83             .expect("overflow")
  84     }
  85
  86     fn with_producer<CB>(self, callback: CB) -> CB::Output
  87         where CB: ProducerCallback<Self::Item>
  88     {
  89         let a_len = self.a.len();
  90         return self.a.with_producer(CallbackA {
  91                                         callback: callback,
  92                                         a_len: a_len,
  93                                         b: self.b,
  94                                     });
  95
  96         struct CallbackA<CB, B> {
  97             callback: CB,
  98             a_len: usize,
  99             b: B,
 100         }
 101
 102         impl<CB, B> ProducerCallback<B::Item> for CallbackA<CB, B>
 103             where B: IndexedParallelIterator,
 104                   CB: ProducerCallback<B::Item>
 105         {
 106             type Output = CB::Output;
 107
 108             fn callback<A>(self, a_producer: A) -> Self::Output
 109                 where A: Producer<Item = B::Item>
 110             {
 111                 return self.b.with_producer(CallbackB {
 112                                                 callback: self.callback,
 113                                                 a_len: self.a_len,
 114                                                 a_producer: a_producer,
 115                                             });
 116             }
 117         }
 118
 119         struct CallbackB<CB, A> {
 120             callback: CB,
 121             a_len: usize,
 122             a_producer: A,
 123         }
 124
 125         impl<CB, A> ProducerCallback<A::Item> for CallbackB<CB, A>
 126             where A: Producer,
 127                   CB: ProducerCallback<A::Item>
 128         {
 129             type Output = CB::Output;
 130
 131             fn callback<B>(self, b_producer: B) -> Self::Output
 132                 where B: Producer<Item = A::Item>
 133             {
 134                 let producer = ChainProducer::new(self.a_len, self.a_producer, b_producer);
 135                 self.callback.callback(producer)
 136             }
 137         }
 138
 139     }
 140 }
 141
 142 /// ////////////////////////////////////////////////////////////////////////
 143
 144 struct ChainProducer<A, B>
 145     where A: Producer,
 146           B: Producer<Item = A::Item>
 147 {
 148     a_len: usize,
 149     a: A,
 150     b: B,
 151 }
 152
 153 impl<A, B> ChainProducer<A, B>
 154     where A: Producer,
 155           B: Producer<Item = A::Item>
 156 {
 157     fn new(a_len: usize, a: A, b: B) -> Self {
 158         ChainProducer {
 159             a_len: a_len,
 160             a: a,
 161             b: b,
 162         }
 163     }
 164 }
 165
 166 impl<A, B> Producer for ChainProducer<A, B>
 167     where A: Producer,
 168           B: Producer<Item = A::Item>
 169 {
 170     type Item = A::Item;
 171     type IntoIter = ChainSeq<A::IntoIter, B::IntoIter>;
 172
 173     fn into_iter(self) -> Self::IntoIter {
 174         ChainSeq::new(self.a.into_iter(), self.b.into_iter())
 175     }
 176
 177     fn min_len(&self) -> usize {
 178         cmp::max(self.a.min_len(), self.b.min_len())
 179     }
 180
 181     fn max_len(&self) -> usize {
 182         cmp::min(self.a.max_len(), self.b.max_len())
 183     }
 184
 185     fn split_at(self, index: usize) -> (Self, Self) {
 186         if index <= self.a_len {
 187             let a_rem = self.a_len - index;
 188             let (a_left, a_right) = self.a.split_at(index);
 189             let (b_left, b_right) = self.b.split_at(0);
 190             (ChainProducer::new(index, a_left, b_left), ChainProducer::new(a_rem, a_right, b_right))
 191         } else {
 192             let (a_left, a_right) = self.a.split_at(self.a_len);
 193             let (b_left, b_right) = self.b.split_at(index - self.a_len);
 194             (ChainProducer::new(self.a_len, a_left, b_left),
 195              ChainProducer::new(0, a_right, b_right))
 196         }
 197     }
 198
 199     fn fold_with<F>(self, mut folder: F) -> F
 200         where F: Folder<A::Item>
 201     {
 202         folder = self.a.fold_with(folder);
 203         if folder.full() {
 204             folder
 205         } else {
 206             self.b.fold_with(folder)
 207         }
 208     }
 209 }
 210
 211 /// ////////////////////////////////////////////////////////////////////////
 212 /// Wrapper for Chain to implement ExactSizeIterator
 213
 214 struct ChainSeq<A, B> {
 215     chain: iter::Chain<A, B>,
 216 }
 217
 218 impl<A, B> ChainSeq<A, B> {
 219     fn new(a: A, b: B) -> ChainSeq<A, B>
 220         where A: ExactSizeIterator,
 221               B: ExactSizeIterator<Item = A::Item>
 222     {
 223         ChainSeq { chain: a.chain(b) }
 224     }
 225 }
 226
 227 impl<A, B> Iterator for ChainSeq<A, B>
 228     where A: Iterator,
 229           B: Iterator<Item = A::Item>
 230 {
 231     type Item = A::Item;
 232
 233     fn next(&mut self) -> Option<Self::Item> {
 234         self.chain.next()
 235     }
 236
 237     fn size_hint(&self) -> (usize, Option<usize>) {
 238         self.chain.size_hint()
 239     }
 240 }
 241
 242 impl<A, B> ExactSizeIterator for ChainSeq<A, B>
 243     where A: ExactSizeIterator,
 244           B: ExactSizeIterator<Item = A::Item>
 245 {
 246 }
 247
 248 impl<A, B> DoubleEndedIterator for ChainSeq<A, B>
 249     where A: DoubleEndedIterator,
 250           B: DoubleEndedIterator<Item = A::Item>
 251 {
 252     fn next_back(&mut self) -> Option<Self::Item> {
 253         self.chain.next_back()
 254     }
 255 }