use super::plumbing::*; use super::*; use std::cmp; use std::iter::Fuse; /// `Interleave` is an iterator that interleaves elements of iterators /// `i` and `j` in one continuous iterator. This struct is created by /// the [`interleave()`] method on [`IndexedParallelIterator`] /// /// [`interleave()`]: trait.IndexedParallelIterator.html#method.interleave /// [`IndexedParallelIterator`]: trait.IndexedParallelIterator.html #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] #[derive(Debug, Clone)] pub struct Interleave where I: IndexedParallelIterator, J: IndexedParallelIterator, { i: I, j: J, } impl Interleave where I: IndexedParallelIterator, J: IndexedParallelIterator, { /// Creates a new `Interleave` iterator pub(super) fn new(i: I, j: J) -> Self { Interleave { i, j } } } impl ParallelIterator for Interleave where I: IndexedParallelIterator, J: IndexedParallelIterator, { type Item = I::Item; fn drive_unindexed(self, consumer: C) -> C::Result where C: Consumer, { bridge(self, consumer) } fn opt_len(&self) -> Option { Some(self.len()) } } impl IndexedParallelIterator for Interleave where I: IndexedParallelIterator, J: IndexedParallelIterator, { fn drive(self, consumer: C) -> C::Result where C: Consumer, { bridge(self, consumer) } fn len(&self) -> usize { self.i.len().checked_add(self.j.len()).expect("overflow") } fn with_producer(self, callback: CB) -> CB::Output where CB: ProducerCallback, { let (i_len, j_len) = (self.i.len(), self.j.len()); return self.i.with_producer(CallbackI { callback, i_len, j_len, i_next: false, j: self.j, }); struct CallbackI { callback: CB, i_len: usize, j_len: usize, i_next: bool, j: J, } impl ProducerCallback for CallbackI where J: IndexedParallelIterator, CB: ProducerCallback, { type Output = CB::Output; fn callback(self, i_producer: I) -> Self::Output where I: Producer, { self.j.with_producer(CallbackJ { i_producer, i_len: self.i_len, j_len: self.j_len, i_next: self.i_next, callback: self.callback, }) } } struct CallbackJ { callback: CB, i_len: usize, j_len: usize, i_next: bool, i_producer: I, } impl ProducerCallback for CallbackJ where I: Producer, CB: ProducerCallback, { type Output = CB::Output; fn callback(self, j_producer: J) -> Self::Output where J: Producer, { let producer = InterleaveProducer::new( self.i_producer, j_producer, self.i_len, self.j_len, self.i_next, ); self.callback.callback(producer) } } } } struct InterleaveProducer where I: Producer, J: Producer, { i: I, j: J, i_len: usize, j_len: usize, i_next: bool, } impl InterleaveProducer where I: Producer, J: Producer, { fn new(i: I, j: J, i_len: usize, j_len: usize, i_next: bool) -> InterleaveProducer { InterleaveProducer { i, j, i_len, j_len, i_next, } } } impl Producer for InterleaveProducer where I: Producer, J: Producer, { type Item = I::Item; type IntoIter = InterleaveSeq; fn into_iter(self) -> Self::IntoIter { InterleaveSeq { i: self.i.into_iter().fuse(), j: self.j.into_iter().fuse(), i_next: self.i_next, } } fn min_len(&self) -> usize { cmp::max(self.i.min_len(), self.j.min_len()) } fn max_len(&self) -> usize { cmp::min(self.i.max_len(), self.j.max_len()) } /// We know 0 < index <= self.i_len + self.j_len /// /// Find a, b satisfying: /// /// (1) 0 < a <= self.i_len /// (2) 0 < b <= self.j_len /// (3) a + b == index /// /// For even splits, set a = b = index/2. /// For odd splits, set a = (index/2)+1, b = index/2, if `i` /// should yield the next element, otherwise, if `j` should yield /// the next element, set a = index/2 and b = (index/2)+1 fn split_at(self, index: usize) -> (Self, Self) { #[inline] fn odd_offset(flag: bool) -> usize { (!flag) as usize } let even = index % 2 == 0; let idx = index >> 1; // desired split let (i_idx, j_idx) = ( idx + odd_offset(even || self.i_next), idx + odd_offset(even || !self.i_next), ); let (i_split, j_split) = if self.i_len >= i_idx && self.j_len >= j_idx { (i_idx, j_idx) } else if self.i_len >= i_idx { // j too short (index - self.j_len, self.j_len) } else { // i too short (self.i_len, index - self.i_len) }; let trailing_i_next = even == self.i_next; let (i_left, i_right) = self.i.split_at(i_split); let (j_left, j_right) = self.j.split_at(j_split); ( InterleaveProducer::new(i_left, j_left, i_split, j_split, self.i_next), InterleaveProducer::new( i_right, j_right, self.i_len - i_split, self.j_len - j_split, trailing_i_next, ), ) } } /// Wrapper for Interleave to implement DoubleEndedIterator and /// ExactSizeIterator. /// /// This iterator is fused. struct InterleaveSeq { i: Fuse, j: Fuse, /// Flag to control which iterator should provide the next element. When /// `false` then `i` produces the next element, otherwise `j` produces the /// next element. i_next: bool, } /// Iterator implementation for InterleaveSeq. This implementation is /// taken more or less verbatim from itertools. It is replicated here /// (instead of calling itertools directly), because we also need to /// implement `DoubledEndedIterator` and `ExactSizeIterator`. impl Iterator for InterleaveSeq where I: Iterator, J: Iterator, { type Item = I::Item; #[inline] fn next(&mut self) -> Option { self.i_next = !self.i_next; if self.i_next { match self.i.next() { None => self.j.next(), r => r, } } else { match self.j.next() { None => self.i.next(), r => r, } } } fn size_hint(&self) -> (usize, Option) { let (ih, jh) = (self.i.size_hint(), self.j.size_hint()); let min = ih.0.saturating_add(jh.0); let max = match (ih.1, jh.1) { (Some(x), Some(y)) => x.checked_add(y), _ => None, }; (min, max) } } // The implementation for DoubleEndedIterator requires // ExactSizeIterator to provide `next_back()`. The last element will // come from the iterator that runs out last (ie has the most elements // in it). If the iterators have the same number of elements, then the // last iterator will provide the last element. impl DoubleEndedIterator for InterleaveSeq where I: DoubleEndedIterator + ExactSizeIterator, J: DoubleEndedIterator + ExactSizeIterator, { #[inline] fn next_back(&mut self) -> Option { match self.i.len().cmp(&self.j.len()) { Ordering::Less => self.j.next_back(), Ordering::Equal => { if self.i_next { self.i.next_back() } else { self.j.next_back() } } Ordering::Greater => self.i.next_back(), } } } impl ExactSizeIterator for InterleaveSeq where I: ExactSizeIterator, J: ExactSizeIterator, { #[inline] fn len(&self) -> usize { self.i.len() + self.j.len() } }