1use super::plumbing::*;
2use super::*;
3
4/// `MultiZip` is an iterator that zips up a tuple of parallel iterators to
5/// produce tuples of their items.
6///
7/// It is created by calling `into_par_iter()` on a tuple of types that
8/// implement `IntoParallelIterator`, or `par_iter()`/`par_iter_mut()` with
9/// types that are iterable by reference.
10///
11/// The implementation currently support tuples up to length 12.
12///
13/// # Examples
14///
15/// ```
16/// use rayon::prelude::*;
17///
18/// // This will iterate `r` by mutable reference, like `par_iter_mut()`, while
19/// // ranges are all iterated by value like `into_par_iter()`.
20/// // Note that the zipped iterator is only as long as the shortest input.
21/// let mut r = vec![0; 3];
22/// (&mut r, 1..10, 10..100, 100..1000).into_par_iter()
23/// .for_each(|(r, x, y, z)| *r = x * y + z);
24///
25/// assert_eq!(&r, &[1 * 10 + 100, 2 * 11 + 101, 3 * 12 + 102]);
26/// ```
27///
28/// For a group that should all be iterated by reference, you can use a tuple reference.
29///
30/// ```
31/// use rayon::prelude::*;
32///
33/// let xs: Vec<_> = (1..10).collect();
34/// let ys: Vec<_> = (10..100).collect();
35/// let zs: Vec<_> = (100..1000).collect();
36///
37/// // Reference each input separately with `IntoParallelIterator`:
38/// let r1: Vec<_> = (&xs, &ys, &zs).into_par_iter()
39/// .map(|(x, y, z)| x * y + z)
40/// .collect();
41///
42/// // Reference them all together with `IntoParallelRefIterator`:
43/// let r2: Vec<_> = (xs, ys, zs).par_iter()
44/// .map(|(x, y, z)| x * y + z)
45/// .collect();
46///
47/// assert_eq!(r1, r2);
48/// ```
49///
50/// Mutable references to a tuple will work similarly.
51///
52/// ```
53/// use rayon::prelude::*;
54///
55/// let mut xs: Vec<_> = (1..4).collect();
56/// let mut ys: Vec<_> = (-4..-1).collect();
57/// let mut zs = vec![0; 3];
58///
59/// // Mutably reference each input separately with `IntoParallelIterator`:
60/// (&mut xs, &mut ys, &mut zs).into_par_iter().for_each(|(x, y, z)| {
61/// *z += *x + *y;
62/// std::mem::swap(x, y);
63/// });
64///
65/// assert_eq!(xs, (vec![-4, -3, -2]));
66/// assert_eq!(ys, (vec![1, 2, 3]));
67/// assert_eq!(zs, (vec![-3, -1, 1]));
68///
69/// // Mutably reference them all together with `IntoParallelRefMutIterator`:
70/// let mut tuple = (xs, ys, zs);
71/// tuple.par_iter_mut().for_each(|(x, y, z)| {
72/// *z += *x + *y;
73/// std::mem::swap(x, y);
74/// });
75///
76/// assert_eq!(tuple, (vec![1, 2, 3], vec![-4, -3, -2], vec![-6, -2, 2]));
77/// ```
78#[derive(Debug, Clone)]
79pub struct MultiZip<T> {
80 tuple: T,
81}
82
83// These macros greedily consume 4 or 2 items first to achieve log2 nesting depth.
84// For example, 5 => 4,1 => (2,2),1.
85//
86// The tuples go up to 12, so we might want to greedily consume 8 too, but
87// the depth works out the same if we let that expand on the right:
88// 9 => 4,5 => (2,2),(4,1) => (2,2),((2,2),1)
89// 12 => 4,8 => (2,2),(4,4) => (2,2),((2,2),(2,2))
90//
91// But if we ever increase to 13, we would want to split 8,5 rather than 4,9.
92
93macro_rules! reduce {
94 ($a:expr, $b:expr, $c:expr, $d:expr, $( $x:expr ),+ => $fn:path) => {
95 reduce!(reduce!($a, $b, $c, $d => $fn),
96 reduce!($( $x ),+ => $fn)
97 => $fn)
98 };
99 ($a:expr, $b:expr, $( $x:expr ),+ => $fn:path) => {
100 reduce!(reduce!($a, $b => $fn),
101 reduce!($( $x ),+ => $fn)
102 => $fn)
103 };
104 ($a:expr, $b:expr => $fn:path) => { $fn($a, $b) };
105 ($a:expr => $fn:path) => { $a };
106}
107
108macro_rules! nest {
109 ($A:tt, $B:tt, $C:tt, $D:tt, $( $X:tt ),+) => {
110 (nest!($A, $B, $C, $D), nest!($( $X ),+))
111 };
112 ($A:tt, $B:tt, $( $X:tt ),+) => {
113 (($A, $B), nest!($( $X ),+))
114 };
115 ($A:tt, $B:tt) => { ($A, $B) };
116 ($A:tt) => { $A };
117}
118
119macro_rules! flatten {
120 ($( $T:ident ),+) => {{
121 #[allow(non_snake_case)]
122 fn flatten<$( $T ),+>(nest!($( $T ),+) : nest!($( $T ),+)) -> ($( $T, )+) {
123 ($( $T, )+)
124 }
125 flatten
126 }};
127}
128
129macro_rules! multizip_impls {
130 ($(
131 $Tuple:ident {
132 $(($idx:tt) -> $T:ident)+
133 }
134 )+) => {
135 $(
136 impl<$( $T, )+> IntoParallelIterator for ($( $T, )+)
137 where
138 $(
139 $T: IntoParallelIterator,
140 $T::Iter: IndexedParallelIterator,
141 )+
142 {
143 type Item = ($( $T::Item, )+);
144 type Iter = MultiZip<($( $T::Iter, )+)>;
145
146 fn into_par_iter(self) -> Self::Iter {
147 MultiZip {
148 tuple: ( $( self.$idx.into_par_iter(), )+ ),
149 }
150 }
151 }
152
153 impl<'a, $( $T, )+> IntoParallelIterator for &'a ($( $T, )+)
154 where
155 $(
156 $T: IntoParallelRefIterator<'a>,
157 $T::Iter: IndexedParallelIterator,
158 )+
159 {
160 type Item = ($( $T::Item, )+);
161 type Iter = MultiZip<($( $T::Iter, )+)>;
162
163 fn into_par_iter(self) -> Self::Iter {
164 MultiZip {
165 tuple: ( $( self.$idx.par_iter(), )+ ),
166 }
167 }
168 }
169
170 impl<'a, $( $T, )+> IntoParallelIterator for &'a mut ($( $T, )+)
171 where
172 $(
173 $T: IntoParallelRefMutIterator<'a>,
174 $T::Iter: IndexedParallelIterator,
175 )+
176 {
177 type Item = ($( $T::Item, )+);
178 type Iter = MultiZip<($( $T::Iter, )+)>;
179
180 fn into_par_iter(self) -> Self::Iter {
181 MultiZip {
182 tuple: ( $( self.$idx.par_iter_mut(), )+ ),
183 }
184 }
185 }
186
187 impl<$( $T, )+> ParallelIterator for MultiZip<($( $T, )+)>
188 where
189 $( $T: IndexedParallelIterator, )+
190 {
191 type Item = ($( $T::Item, )+);
192
193 fn drive_unindexed<CONSUMER>(self, consumer: CONSUMER) -> CONSUMER::Result
194 where
195 CONSUMER: UnindexedConsumer<Self::Item>,
196 {
197 self.drive(consumer)
198 }
199
200 fn opt_len(&self) -> Option<usize> {
201 Some(self.len())
202 }
203 }
204
205 impl<$( $T, )+> IndexedParallelIterator for MultiZip<($( $T, )+)>
206 where
207 $( $T: IndexedParallelIterator, )+
208 {
209 fn drive<CONSUMER>(self, consumer: CONSUMER) -> CONSUMER::Result
210 where
211 CONSUMER: Consumer<Self::Item>,
212 {
213 reduce!($( self.tuple.$idx ),+ => IndexedParallelIterator::zip)
214 .map(flatten!($( $T ),+))
215 .drive(consumer)
216 }
217
218 fn len(&self) -> usize {
219 reduce!($( self.tuple.$idx.len() ),+ => Ord::min)
220 }
221
222 fn with_producer<CB>(self, callback: CB) -> CB::Output
223 where
224 CB: ProducerCallback<Self::Item>,
225 {
226 reduce!($( self.tuple.$idx ),+ => IndexedParallelIterator::zip)
227 .map(flatten!($( $T ),+))
228 .with_producer(callback)
229 }
230 }
231 )+
232 }
233}
234
235multizip_impls! {
236 Tuple1 {
237 (0) -> A
238 }
239 Tuple2 {
240 (0) -> A
241 (1) -> B
242 }
243 Tuple3 {
244 (0) -> A
245 (1) -> B
246 (2) -> C
247 }
248 Tuple4 {
249 (0) -> A
250 (1) -> B
251 (2) -> C
252 (3) -> D
253 }
254 Tuple5 {
255 (0) -> A
256 (1) -> B
257 (2) -> C
258 (3) -> D
259 (4) -> E
260 }
261 Tuple6 {
262 (0) -> A
263 (1) -> B
264 (2) -> C
265 (3) -> D
266 (4) -> E
267 (5) -> F
268 }
269 Tuple7 {
270 (0) -> A
271 (1) -> B
272 (2) -> C
273 (3) -> D
274 (4) -> E
275 (5) -> F
276 (6) -> G
277 }
278 Tuple8 {
279 (0) -> A
280 (1) -> B
281 (2) -> C
282 (3) -> D
283 (4) -> E
284 (5) -> F
285 (6) -> G
286 (7) -> H
287 }
288 Tuple9 {
289 (0) -> A
290 (1) -> B
291 (2) -> C
292 (3) -> D
293 (4) -> E
294 (5) -> F
295 (6) -> G
296 (7) -> H
297 (8) -> I
298 }
299 Tuple10 {
300 (0) -> A
301 (1) -> B
302 (2) -> C
303 (3) -> D
304 (4) -> E
305 (5) -> F
306 (6) -> G
307 (7) -> H
308 (8) -> I
309 (9) -> J
310 }
311 Tuple11 {
312 (0) -> A
313 (1) -> B
314 (2) -> C
315 (3) -> D
316 (4) -> E
317 (5) -> F
318 (6) -> G
319 (7) -> H
320 (8) -> I
321 (9) -> J
322 (10) -> K
323 }
324 Tuple12 {
325 (0) -> A
326 (1) -> B
327 (2) -> C
328 (3) -> D
329 (4) -> E
330 (5) -> F
331 (6) -> G
332 (7) -> H
333 (8) -> I
334 (9) -> J
335 (10) -> K
336 (11) -> L
337 }
338}
339