]>
Commit | Line | Data |
---|---|---|
6a06907d XL |
1 | use super::plumbing::*; |
2 | use 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)] | |
79 | pub 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 | ||
93 | macro_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 | ||
108 | macro_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 | ||
119 | macro_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 | ||
129 | macro_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 | ||
235 | multizip_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 | } |