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

   1 use super::plumbing::*;
   2 use super::*;
   3 use std::cmp;
   4 use std::iter::Fuse;
   5
   6 /// `Interleave` is an iterator that interleaves elements of iterators
   7 /// `i` and `j` in one continuous iterator. This struct is created by
   8 /// the [`interleave()`] method on [`IndexedParallelIterator`]
   9 ///
  10 /// [`interleave()`]: trait.IndexedParallelIterator.html#method.interleave
  11 /// [`IndexedParallelIterator`]: trait.IndexedParallelIterator.html
  12 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
  13 #[derive(Debug, Clone)]
  14 pub struct Interleave<I, J>
  15 where
  16     I: IndexedParallelIterator,
  17     J: IndexedParallelIterator<Item = I::Item>,
  18 {
  19     i: I,
  20     j: J,
  21 }
  22
  23 impl<I, J> Interleave<I, J>
  24 where
  25     I: IndexedParallelIterator,
  26     J: IndexedParallelIterator<Item = I::Item>,
  27 {
  28     /// Creates a new `Interleave` iterator
  29     pub(super) fn new(i: I, j: J) -> Self {
  30         Interleave { i, j }
  31     }
  32 }
  33
  34 impl<I, J> ParallelIterator for Interleave<I, J>
  35 where
  36     I: IndexedParallelIterator,
  37     J: IndexedParallelIterator<Item = I::Item>,
  38 {
  39     type Item = I::Item;
  40
  41     fn drive_unindexed<C>(self, consumer: C) -> C::Result
  42     where
  43         C: Consumer<I::Item>,
  44     {
  45         bridge(self, consumer)
  46     }
  47
  48     fn opt_len(&self) -> Option<usize> {
  49         Some(self.len())
  50     }
  51 }
  52
  53 impl<I, J> IndexedParallelIterator for Interleave<I, J>
  54 where
  55     I: IndexedParallelIterator,
  56     J: IndexedParallelIterator<Item = I::Item>,
  57 {
  58     fn drive<C>(self, consumer: C) -> C::Result
  59     where
  60         C: Consumer<Self::Item>,
  61     {
  62         bridge(self, consumer)
  63     }
  64
  65     fn len(&self) -> usize {
  66         self.i.len().checked_add(self.j.len()).expect("overflow")
  67     }
  68
  69     fn with_producer<CB>(self, callback: CB) -> CB::Output
  70     where
  71         CB: ProducerCallback<Self::Item>,
  72     {
  73         let (i_len, j_len) = (self.i.len(), self.j.len());
  74         return self.i.with_producer(CallbackI {
  75             callback,
  76             i_len,
  77             j_len,
  78             i_next: false,
  79             j: self.j,
  80         });
  81
  82         struct CallbackI<CB, J> {
  83             callback: CB,
  84             i_len: usize,
  85             j_len: usize,
  86             i_next: bool,
  87             j: J,
  88         }
  89
  90         impl<CB, J> ProducerCallback<J::Item> for CallbackI<CB, J>
  91         where
  92             J: IndexedParallelIterator,
  93             CB: ProducerCallback<J::Item>,
  94         {
  95             type Output = CB::Output;
  96
  97             fn callback<I>(self, i_producer: I) -> Self::Output
  98             where
  99                 I: Producer<Item = J::Item>,
 100             {
 101                 self.j.with_producer(CallbackJ {
 102                     i_producer,
 103                     i_len: self.i_len,
 104                     j_len: self.j_len,
 105                     i_next: self.i_next,
 106                     callback: self.callback,
 107                 })
 108             }
 109         }
 110
 111         struct CallbackJ<CB, I> {
 112             callback: CB,
 113             i_len: usize,
 114             j_len: usize,
 115             i_next: bool,
 116             i_producer: I,
 117         }
 118
 119         impl<CB, I> ProducerCallback<I::Item> for CallbackJ<CB, I>
 120         where
 121             I: Producer,
 122             CB: ProducerCallback<I::Item>,
 123         {
 124             type Output = CB::Output;
 125
 126             fn callback<J>(self, j_producer: J) -> Self::Output
 127             where
 128                 J: Producer<Item = I::Item>,
 129             {
 130                 let producer = InterleaveProducer::new(
 131                     self.i_producer,
 132                     j_producer,
 133                     self.i_len,
 134                     self.j_len,
 135                     self.i_next,
 136                 );
 137                 self.callback.callback(producer)
 138             }
 139         }
 140     }
 141 }
 142
 143 struct InterleaveProducer<I, J>
 144 where
 145     I: Producer,
 146     J: Producer<Item = I::Item>,
 147 {
 148     i: I,
 149     j: J,
 150     i_len: usize,
 151     j_len: usize,
 152     i_next: bool,
 153 }
 154
 155 impl<I, J> InterleaveProducer<I, J>
 156 where
 157     I: Producer,
 158     J: Producer<Item = I::Item>,
 159 {
 160     fn new(i: I, j: J, i_len: usize, j_len: usize, i_next: bool) -> InterleaveProducer<I, J> {
 161         InterleaveProducer {
 162             i,
 163             j,
 164             i_len,
 165             j_len,
 166             i_next,
 167         }
 168     }
 169 }
 170
 171 impl<I, J> Producer for InterleaveProducer<I, J>
 172 where
 173     I: Producer,
 174     J: Producer<Item = I::Item>,
 175 {
 176     type Item = I::Item;
 177     type IntoIter = InterleaveSeq<I::IntoIter, J::IntoIter>;
 178
 179     fn into_iter(self) -> Self::IntoIter {
 180         InterleaveSeq {
 181             i: self.i.into_iter().fuse(),
 182             j: self.j.into_iter().fuse(),
 183             i_next: self.i_next,
 184         }
 185     }
 186
 187     fn min_len(&self) -> usize {
 188         cmp::max(self.i.min_len(), self.j.min_len())
 189     }
 190
 191     fn max_len(&self) -> usize {
 192         cmp::min(self.i.max_len(), self.j.max_len())
 193     }
 194
 195     /// We know 0 < index <= self.i_len + self.j_len
 196     ///
 197     /// Find a, b satisfying:
 198     ///
 199     ///  (1) 0 < a <= self.i_len
 200     ///  (2) 0 < b <= self.j_len
 201     ///  (3) a + b == index
 202     ///
 203     /// For even splits, set a = b = index/2.
 204     /// For odd splits, set a = (index/2)+1, b = index/2, if `i`
 205     /// should yield the next element, otherwise, if `j` should yield
 206     /// the next element, set a = index/2 and b = (index/2)+1
 207     fn split_at(self, index: usize) -> (Self, Self) {
 208         #[inline]
 209         fn odd_offset(flag: bool) -> usize {
 210             (!flag) as usize
 211         }
 212
 213         let even = index % 2 == 0;
 214         let idx = index >> 1;
 215
 216         // desired split
 217         let (i_idx, j_idx) = (
 218             idx + odd_offset(even || self.i_next),
 219             idx + odd_offset(even || !self.i_next),
 220         );
 221
 222         let (i_split, j_split) = if self.i_len >= i_idx && self.j_len >= j_idx {
 223             (i_idx, j_idx)
 224         } else if self.i_len >= i_idx {
 225             // j too short
 226             (index - self.j_len, self.j_len)
 227         } else {
 228             // i too short
 229             (self.i_len, index - self.i_len)
 230         };
 231
 232         let trailing_i_next = even == self.i_next;
 233         let (i_left, i_right) = self.i.split_at(i_split);
 234         let (j_left, j_right) = self.j.split_at(j_split);
 235
 236         (
 237             InterleaveProducer::new(i_left, j_left, i_split, j_split, self.i_next),
 238             InterleaveProducer::new(
 239                 i_right,
 240                 j_right,
 241                 self.i_len - i_split,
 242                 self.j_len - j_split,
 243                 trailing_i_next,
 244             ),
 245         )
 246     }
 247 }
 248
 249 /// Wrapper for Interleave to implement DoubleEndedIterator and
 250 /// ExactSizeIterator.
 251 ///
 252 /// This iterator is fused.
 253 struct InterleaveSeq<I, J> {
 254     i: Fuse<I>,
 255     j: Fuse<J>,
 256
 257     /// Flag to control which iterator should provide the next element. When
 258     /// `false` then `i` produces the next element, otherwise `j` produces the
 259     /// next element.
 260     i_next: bool,
 261 }
 262
 263 /// Iterator implementation for InterleaveSeq. This implementation is
 264 /// taken more or less verbatim from itertools. It is replicated here
 265 /// (instead of calling itertools directly), because we also need to
 266 /// implement `DoubledEndedIterator` and `ExactSizeIterator`.
 267 impl<I, J> Iterator for InterleaveSeq<I, J>
 268 where
 269     I: Iterator,
 270     J: Iterator<Item = I::Item>,
 271 {
 272     type Item = I::Item;
 273
 274     #[inline]
 275     fn next(&mut self) -> Option<Self::Item> {
 276         self.i_next = !self.i_next;
 277         if self.i_next {
 278             match self.i.next() {
 279                 None => self.j.next(),
 280                 r => r,
 281             }
 282         } else {
 283             match self.j.next() {
 284                 None => self.i.next(),
 285                 r => r,
 286             }
 287         }
 288     }
 289
 290     fn size_hint(&self) -> (usize, Option<usize>) {
 291         let (ih, jh) = (self.i.size_hint(), self.j.size_hint());
 292         let min = ih.0.saturating_add(jh.0);
 293         let max = match (ih.1, jh.1) {
 294             (Some(x), Some(y)) => x.checked_add(y),
 295             _ => None,
 296         };
 297         (min, max)
 298     }
 299 }
 300
 301 // The implementation for DoubleEndedIterator requires
 302 // ExactSizeIterator to provide `next_back()`. The last element will
 303 // come from the iterator that runs out last (ie has the most elements
 304 // in it). If the iterators have the same number of elements, then the
 305 // last iterator will provide the last element.
 306 impl<I, J> DoubleEndedIterator for InterleaveSeq<I, J>
 307 where
 308     I: DoubleEndedIterator + ExactSizeIterator,
 309     J: DoubleEndedIterator<Item = I::Item> + ExactSizeIterator<Item = I::Item>,
 310 {
 311     #[inline]
 312     fn next_back(&mut self) -> Option<I::Item> {
 313         if self.i.len() == self.j.len() {
 314             if self.i_next {
 315                 self.i.next_back()
 316             } else {
 317                 self.j.next_back()
 318             }
 319         } else if self.i.len() < self.j.len() {
 320             self.j.next_back()
 321         } else {
 322             self.i.next_back()
 323         }
 324     }
 325 }
 326
 327 impl<I, J> ExactSizeIterator for InterleaveSeq<I, J>
 328 where
 329     I: ExactSizeIterator,
 330     J: ExactSizeIterator<Item = I::Item>,
 331 {
 332     #[inline]
 333     fn len(&self) -> usize {
 334         self.i.len() + self.j.len()
 335     }
 336 }