Skip to main content

moirai_iter/parallel/adapters/
pair.rs

1use super::super::{Consumer, ParallelIterator, VecParIter};
2
3fn interleave_all<T>(left: Vec<T>, right: Vec<T>) -> Vec<T> {
4    let len = left.len().checked_add(right.len()).expect("overflow");
5    let paired_len = left.len().min(right.len());
6    let mut left = left.into_iter();
7    let mut right = right.into_iter();
8    let mut output = Vec::with_capacity(len);
9
10    for _ in 0..paired_len {
11        output.push(left.next().expect("paired length is bounded by left input"));
12        output.push(
13            right
14                .next()
15                .expect("paired length is bounded by right input"),
16        );
17    }
18    output.extend(left);
19    output.extend(right);
20    output
21}
22
23fn interleave_shortest<T>(mut left: Vec<T>, mut right: Vec<T>) -> Vec<T> {
24    let left_take = if left.len() <= right.len() {
25        left.len()
26    } else {
27        right.len() + 1
28    };
29    let right_take = right.len().min(left.len());
30
31    let output_len = left_take.checked_add(right_take).expect("overflow");
32    left.truncate(left_take);
33    right.truncate(right_take);
34
35    let mut left = left.into_iter();
36    let mut right = right.into_iter();
37    let mut output = Vec::with_capacity(output_len);
38    for _ in 0..right_take {
39        output.push(left.next().expect("take count is bounded by left input"));
40        output.push(right.next().expect("take count is bounded by right input"));
41    }
42    output.extend(left);
43    output
44}
45
46/// Zip adapter with standard shortest-input value semantics.
47pub struct Zip<I, J> {
48    left: I,
49    right: J,
50}
51
52impl<I, J> Zip<I, J> {
53    pub(in crate::parallel) fn new(left: I, right: J) -> Self {
54        Self { left, right }
55    }
56}
57
58impl<I, J> ParallelIterator for Zip<I, J>
59where
60    I: ParallelIterator,
61    J: ParallelIterator,
62    I::Item: Sync + 'static,
63    J::Item: Sync + 'static,
64{
65    type Item = (I::Item, J::Item);
66
67    fn seq_items(self) -> Vec<Self::Item> {
68        self.left
69            .seq_items()
70            .into_iter()
71            .zip(self.right.seq_items())
72            .collect()
73    }
74
75    fn drive<C, R>(self, consumer: C) -> R
76    where
77        C: Consumer<Self::Item, Result = R> + Send + Sync,
78        R: Send,
79    {
80        consumer.consume(VecParIter::new(self.seq_items()))
81    }
82}
83
84/// Zip adapter with Rayon-style equal-length value semantics.
85pub struct ZipEq<I, J> {
86    pub(in crate::parallel) left: I,
87    pub(in crate::parallel) right: J,
88}
89
90impl<I, J> ZipEq<I, J> {
91    pub(in crate::parallel) fn new(left: I, right: J) -> Self {
92        Self { left, right }
93    }
94}
95
96impl<I, J> ParallelIterator for ZipEq<I, J>
97where
98    I: ParallelIterator,
99    J: ParallelIterator,
100    I::Item: Sync + 'static,
101    J::Item: Sync + 'static,
102{
103    type Item = (I::Item, J::Item);
104
105    fn seq_items(self) -> Vec<Self::Item> {
106        let left = self.left.seq_items();
107        let right = self.right.seq_items();
108        assert_eq!(
109            left.len(),
110            right.len(),
111            "zip_eq requires equal input lengths"
112        );
113        left.into_iter().zip(right).collect()
114    }
115
116    fn drive<C, R>(self, consumer: C) -> R
117    where
118        C: Consumer<Self::Item, Result = R> + Send + Sync,
119        R: Send,
120    {
121        consumer.consume(VecParIter::new(self.seq_items()))
122    }
123}
124
125/// Indexed interleave adapter with Rayon-style alternating value semantics.
126pub struct Interleave<I, J> {
127    pub(in crate::parallel) left: I,
128    pub(in crate::parallel) right: J,
129}
130
131impl<I, J> Interleave<I, J> {
132    pub(in crate::parallel) fn new(left: I, right: J) -> Self {
133        Self { left, right }
134    }
135}
136
137impl<I, J> ParallelIterator for Interleave<I, J>
138where
139    I: ParallelIterator,
140    J: ParallelIterator<Item = I::Item>,
141    I::Item: Sync + 'static,
142{
143    type Item = I::Item;
144
145    fn seq_items(self) -> Vec<Self::Item> {
146        interleave_all(self.left.seq_items(), self.right.seq_items())
147    }
148
149    fn drive<C, R>(self, consumer: C) -> R
150    where
151        C: Consumer<Self::Item, Result = R> + Send + Sync,
152        R: Send,
153    {
154        consumer.consume(VecParIter::new(self.seq_items()))
155    }
156}
157
158/// Indexed interleave adapter that truncates at the shorter input boundary.
159pub struct InterleaveShortest<I, J> {
160    pub(in crate::parallel) left: I,
161    pub(in crate::parallel) right: J,
162}
163
164impl<I, J> InterleaveShortest<I, J> {
165    pub(in crate::parallel) fn new(left: I, right: J) -> Self {
166        Self { left, right }
167    }
168}
169
170impl<I, J> ParallelIterator for InterleaveShortest<I, J>
171where
172    I: ParallelIterator,
173    J: ParallelIterator<Item = I::Item>,
174    I::Item: Sync + 'static,
175{
176    type Item = I::Item;
177
178    fn seq_items(self) -> Vec<Self::Item> {
179        interleave_shortest(self.left.seq_items(), self.right.seq_items())
180    }
181
182    fn drive<C, R>(self, consumer: C) -> R
183    where
184        C: Consumer<Self::Item, Result = R> + Send + Sync,
185        R: Send,
186    {
187        consumer.consume(VecParIter::new(self.seq_items()))
188    }
189}