Skip to content

Commit 0aefd61

Browse files
bors[bot]cuviper
bors[bot]
and
cuviper
authored
711: impl IntoParallelIterator for tuples => MultiZip r=nikomatsakis a=cuviper This is implemented for tuples up to arity 12, much like the standard library's trait implementations. - For `(a, b, ...)`, it calls `into_par_iter()` on each member. - For `&(a, b, ...)`, it calls `par_iter()` on each member. - For `&mut (a, b, ...)`, it calls `par_iter_mut()` on each member. The resulting `MultiZip` iterator returns a tuple of the zipped items from each input iterator. Internally, it is implemented with macros that forward to a series of regular `zip`s, mapping to a flattened tuple. Closes rayon-rs#567. Co-authored-by: Josh Stone <[email protected]>
2 parents b6449cb + c0dd154 commit 0aefd61

File tree

4 files changed

+374
-0
lines changed

4 files changed

+374
-0
lines changed

‎src/iter/mod.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,8 @@ mod zip;
141141
pub use self::zip::Zip;
142142
mod zip_eq;
143143
pub use self::zip_eq::ZipEq;
144+
mod multizip;
145+
pub use self::multizip::MultiZip;
144146
mod interleave;
145147
pub use self::interleave::Interleave;
146148
mod interleave_shortest;

‎src/iter/multizip.rs

Lines changed: 338 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,338 @@
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+
}

‎tests/clones.rs

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -162,3 +162,20 @@ fn clone_repeat() {
162162
fn clone_splitter() {
163163
check(rayon::iter::split(0..1000, |x| (x, None)));
164164
}
165+
166+
#[test]
167+
fn clone_multizip() {
168+
let v: &Vec<_> = &(0..1000).collect();
169+
check((v,).into_par_iter());
170+
check((v, v).into_par_iter());
171+
check((v, v, v).into_par_iter());
172+
check((v, v, v, v).into_par_iter());
173+
check((v, v, v, v, v).into_par_iter());
174+
check((v, v, v, v, v, v).into_par_iter());
175+
check((v, v, v, v, v, v, v).into_par_iter());
176+
check((v, v, v, v, v, v, v, v).into_par_iter());
177+
check((v, v, v, v, v, v, v, v, v).into_par_iter());
178+
check((v, v, v, v, v, v, v, v, v, v).into_par_iter());
179+
check((v, v, v, v, v, v, v, v, v, v, v).into_par_iter());
180+
check((v, v, v, v, v, v, v, v, v, v, v, v).into_par_iter());
181+
}

‎tests/debug.rs

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -173,3 +173,20 @@ fn debug_repeat() {
173173
fn debug_splitter() {
174174
check(rayon::iter::split(0..10, |x| (x, None)));
175175
}
176+
177+
#[test]
178+
fn debug_multizip() {
179+
let v: &Vec<_> = &(0..10).collect();
180+
check((v,).into_par_iter());
181+
check((v, v).into_par_iter());
182+
check((v, v, v).into_par_iter());
183+
check((v, v, v, v).into_par_iter());
184+
check((v, v, v, v, v).into_par_iter());
185+
check((v, v, v, v, v, v).into_par_iter());
186+
check((v, v, v, v, v, v, v).into_par_iter());
187+
check((v, v, v, v, v, v, v, v).into_par_iter());
188+
check((v, v, v, v, v, v, v, v, v).into_par_iter());
189+
check((v, v, v, v, v, v, v, v, v, v).into_par_iter());
190+
check((v, v, v, v, v, v, v, v, v, v, v).into_par_iter());
191+
check((v, v, v, v, v, v, v, v, v, v, v, v).into_par_iter());
192+
}

0 commit comments

Comments
 (0)