diff options
author | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-04-07 19:33:14 +0000 |
---|---|---|
committer | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-04-07 19:33:14 +0000 |
commit | 36d22d82aa202bb199967e9512281e9a53db42c9 (patch) | |
tree | 105e8c98ddea1c1e4784a60a5a6410fa416be2de /third_party/rust/rayon/src/iter/interleave.rs | |
parent | Initial commit. (diff) | |
download | firefox-esr-upstream.tar.xz firefox-esr-upstream.zip |
Adding upstream version 115.7.0esr.upstream/115.7.0esrupstream
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'third_party/rust/rayon/src/iter/interleave.rs')
-rw-r--r-- | third_party/rust/rayon/src/iter/interleave.rs | 336 |
1 files changed, 336 insertions, 0 deletions
diff --git a/third_party/rust/rayon/src/iter/interleave.rs b/third_party/rust/rayon/src/iter/interleave.rs new file mode 100644 index 0000000000..3cacc49f95 --- /dev/null +++ b/third_party/rust/rayon/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> { + 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<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() + } +} |