aboutsummaryrefslogtreecommitdiff
path: root/src/iter/interleave.rs
diff options
context:
space:
mode:
Diffstat (limited to 'src/iter/interleave.rs')
-rw-r--r--src/iter/interleave.rs336
1 files changed, 336 insertions, 0 deletions
diff --git a/src/iter/interleave.rs b/src/iter/interleave.rs
new file mode 100644
index 0000000..b5d43d5
--- /dev/null
+++ b/src/iter/interleave.rs
@@ -0,0 +1,336 @@
+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<I, J>
+where
+ I: IndexedParallelIterator,
+ J: IndexedParallelIterator<Item = I::Item>,
+{
+ i: I,
+ j: J,
+}
+
+impl<I, J> Interleave<I, J>
+where
+ I: IndexedParallelIterator,
+ J: IndexedParallelIterator<Item = I::Item>,
+{
+ /// Creates a new `Interleave` iterator
+ pub(super) fn new(i: I, j: J) -> Self {
+ Interleave { i, j }
+ }
+}
+
+impl<I, J> ParallelIterator for Interleave<I, J>
+where
+ I: IndexedParallelIterator,
+ J: IndexedParallelIterator<Item = I::Item>,
+{
+ type Item = I::Item;
+
+ fn drive_unindexed<C>(self, consumer: C) -> C::Result
+ where
+ C: Consumer<I::Item>,
+ {
+ bridge(self, consumer)
+ }
+
+ fn opt_len(&self) -> Option<usize> {
+ Some(self.len())
+ }
+}
+
+impl<I, J> IndexedParallelIterator for Interleave<I, J>
+where
+ I: IndexedParallelIterator,
+ J: IndexedParallelIterator<Item = I::Item>,
+{
+ fn drive<C>(self, consumer: C) -> C::Result
+ where
+ C: Consumer<Self::Item>,
+ {
+ bridge(self, consumer)
+ }
+
+ fn len(&self) -> usize {
+ self.i.len().checked_add(self.j.len()).expect("overflow")
+ }
+
+ fn with_producer<CB>(self, callback: CB) -> CB::Output
+ where
+ CB: ProducerCallback<Self::Item>,
+ {
+ 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<CB, J> {
+ callback: CB,
+ i_len: usize,
+ j_len: usize,
+ i_next: bool,
+ j: J,
+ }
+
+ impl<CB, J> ProducerCallback<J::Item> for CallbackI<CB, J>
+ where
+ J: IndexedParallelIterator,
+ CB: ProducerCallback<J::Item>,
+ {
+ type Output = CB::Output;
+
+ fn callback<I>(self, i_producer: I) -> Self::Output
+ where
+ I: Producer<Item = J::Item>,
+ {
+ 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<CB, I> {
+ callback: CB,
+ i_len: usize,
+ j_len: usize,
+ i_next: bool,
+ i_producer: I,
+ }
+
+ impl<CB, I> ProducerCallback<I::Item> for CallbackJ<CB, I>
+ where
+ I: Producer,
+ CB: ProducerCallback<I::Item>,
+ {
+ type Output = CB::Output;
+
+ fn callback<J>(self, j_producer: J) -> Self::Output
+ where
+ J: Producer<Item = I::Item>,
+ {
+ let producer = InterleaveProducer::new(
+ self.i_producer,
+ j_producer,
+ self.i_len,
+ self.j_len,
+ self.i_next,
+ );
+ self.callback.callback(producer)
+ }
+ }
+ }
+}
+
+struct InterleaveProducer<I, J>
+where
+ I: Producer,
+ J: Producer<Item = I::Item>,
+{
+ i: I,
+ j: J,
+ i_len: usize,
+ j_len: usize,
+ i_next: bool,
+}
+
+impl<I, J> InterleaveProducer<I, J>
+where
+ I: Producer,
+ J: Producer<Item = I::Item>,
+{
+ fn new(i: I, j: J, i_len: usize, j_len: usize, i_next: bool) -> InterleaveProducer<I, J> {
+ InterleaveProducer {
+ i,
+ j,
+ i_len,
+ j_len,
+ i_next,
+ }
+ }
+}
+
+impl<I, J> Producer for InterleaveProducer<I, J>
+where
+ I: Producer,
+ J: Producer<Item = I::Item>,
+{
+ type Item = I::Item;
+ type IntoIter = InterleaveSeq<I::IntoIter, J::IntoIter>;
+
+ 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, J> {
+ i: Fuse<I>,
+ j: Fuse<J>,
+
+ /// 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<I, J> Iterator for InterleaveSeq<I, J>
+where
+ I: Iterator,
+ J: Iterator<Item = I::Item>,
+{
+ type Item = I::Item;
+
+ #[inline]
+ fn next(&mut self) -> Option<Self::Item> {
+ 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<usize>) {
+ 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<I, J> DoubleEndedIterator for InterleaveSeq<I, J>
+where
+ I: DoubleEndedIterator + ExactSizeIterator,
+ J: DoubleEndedIterator<Item = I::Item> + ExactSizeIterator<Item = I::Item>,
+{
+ #[inline]
+ fn next_back(&mut self) -> Option<I::Item> {
+ if self.i.len() == self.j.len() {
+ if self.i_next {
+ self.i.next_back()
+ } else {
+ self.j.next_back()
+ }
+ } else if self.i.len() < self.j.len() {
+ self.j.next_back()
+ } else {
+ self.i.next_back()
+ }
+ }
+}
+
+impl<I, J> ExactSizeIterator for InterleaveSeq<I, J>
+where
+ I: ExactSizeIterator,
+ J: ExactSizeIterator<Item = I::Item>,
+{
+ #[inline]
+ fn len(&self) -> usize {
+ self.i.len() + self.j.len()
+ }
+}
generated by cgit v1.2.3 (git 2.25.1) at 2024-07-05 22:13:35 +0000