moirai_iter/parallel/adapters/
pair.rs1use 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
46pub 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
84pub 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
125pub 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
158pub 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}