From faf347a416e8998ace15e795d0af9040cd0ec15e Mon Sep 17 00:00:00 2001 From: Josh Stone Date: Thu, 5 Dec 2019 09:15:00 -0800 Subject: [PATCH 1/3] impl IntoParallelIterator for tuples => MultiZip 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. --- src/iter/mod.rs | 2 + src/iter/multizip.rs | 321 +++++++++++++++++++++++++++++++++++++++++++ tests/clones.rs | 17 +++ tests/debug.rs | 17 +++ 4 files changed, 357 insertions(+) create mode 100644 src/iter/multizip.rs diff --git a/src/iter/mod.rs b/src/iter/mod.rs index 4d91f7dde..21f82f4fd 100644 --- a/src/iter/mod.rs +++ b/src/iter/mod.rs @@ -141,6 +141,8 @@ mod zip; pub use self::zip::Zip; mod zip_eq; pub use self::zip_eq::ZipEq; +mod multizip; +pub use self::multizip::MultiZip; mod interleave; pub use self::interleave::Interleave; mod interleave_shortest; diff --git a/src/iter/multizip.rs b/src/iter/multizip.rs new file mode 100644 index 000000000..bb4f130bb --- /dev/null +++ b/src/iter/multizip.rs @@ -0,0 +1,321 @@ +use super::plumbing::*; +use super::*; + +use std::cmp; + +/// `MultiZip` is an iterator that zips up a tuple of parallel iterators to +/// produce tuples of their items. +/// +/// It is created by calling `into_par_iter()` on a tuple of types that +/// implement `IntoParallelIterator`, or `par_iter()`/`par_iter_mut()` with +/// types that are iterable by reference. +/// +/// The implementation currently support tuples up to length 12. +/// +/// # Examples +/// +/// ``` +/// use rayon::prelude::*; +/// +/// // This will iterate `r` by mutable reference, like `par_iter_mut()`, while +/// // ranges are all iterated by value like `into_par_iter()`. +/// // Note that the zipped iterator is only as long as the shortest input. +/// let mut r = vec![0; 3]; +/// (&mut r, 1..10, 10..100, 100..1000).into_par_iter() +/// .for_each(|(r, x, y, z)| *r = x * y + z); +/// +/// assert_eq!(&r, &[1 * 10 + 100, 2 * 11 + 101, 3 * 12 + 102]); +/// ``` +/// +/// For a group that should all be iterated by reference, you can use a tuple reference. +/// +/// ``` +/// use rayon::prelude::*; +/// +/// let xs: Vec<_> = (1..10).collect(); +/// let ys: Vec<_> = (10..100).collect(); +/// let zs: Vec<_> = (100..1000).collect(); +/// +/// // Reference each input separately with `IntoParallelIterator`: +/// let r1: Vec<_> = (&xs, &ys, &zs).into_par_iter() +/// .map(|(x, y, z)| x * y + z) +/// .collect(); +/// +/// // Reference them all together with `IntoParallelRefIterator`: +/// let r2: Vec<_> = (xs, ys, zs).par_iter() +/// .map(|(x, y, z)| x * y + z) +/// .collect(); +/// +/// assert_eq!(r1, r2); +/// ``` +/// +/// Mutable references to a tuple will work similarly. +/// +/// ``` +/// use rayon::prelude::*; +/// +/// let mut xs: Vec<_> = (1..4).collect(); +/// let mut ys: Vec<_> = (-4..-1).collect(); +/// let mut zs = vec![0; 3]; +/// +/// // Mutably reference each input separately with `IntoParallelIterator`: +/// (&mut xs, &mut ys, &mut zs).into_par_iter().for_each(|(x, y, z)| { +/// *z += *x + *y; +/// std::mem::swap(x, y); +/// }); +/// +/// assert_eq!(xs, (vec![-4, -3, -2])); +/// assert_eq!(ys, (vec![1, 2, 3])); +/// assert_eq!(zs, (vec![-3, -1, 1])); +/// +/// // Mutably reference them all together with `IntoParallelRefMutIterator`: +/// let mut tuple = (xs, ys, zs); +/// tuple.par_iter_mut().for_each(|(x, y, z)| { +/// *z += *x + *y; +/// std::mem::swap(x, y); +/// }); +/// +/// assert_eq!(tuple, (vec![1, 2, 3], vec![-4, -3, -2], vec![-6, -2, 2])); +/// ``` +#[derive(Debug, Clone)] +pub struct MultiZip { + tuple: T, +} + +macro_rules! zip { + ($first:expr, $( $iter:expr, )*) => { + $first $( .zip($iter) )* + }; +} + +macro_rules! min { + ($x:expr,) => { $x }; + ($x:expr, $( $y:expr, )+) => { cmp::min($x, min!($( $y, )+)) }; +} + +macro_rules! flatten { + (|$a:tt : $A:tt| -> ($( $X:ident, )+) { $tuple:tt };) => {{ + fn flatten<$( $X ),+>($a : $A) -> ($( $X, )*) { + $tuple + } + flatten + }}; + (|$a:tt : $A:tt| -> ($( $X:ident, )+) { ($( $x:ident, )+) }; + $B:ident, $( $T:ident, )*) => { + flatten!(|($a, b): ($A, $B)| -> ($( $X, )+ $B,) { ($( $x, )+ b,) }; $( $T, )*) + }; + ($A:ident, $( $T:ident, )*) => { + flatten!(|a: $A| -> ($A,) { (a,) }; $( $T, )*) + }; +} + +macro_rules! multizip_impls { + ($( + $Tuple:ident { + $(($idx:tt) -> $T:ident)+ + } + )+) => { + $( + impl<$( $T, )+> IntoParallelIterator for ($( $T, )+) + where + $( + $T: IntoParallelIterator, + $T::Iter: IndexedParallelIterator, + )+ + { + type Item = ($( $T::Item, )+); + type Iter = MultiZip<($( $T::Iter, )+)>; + + fn into_par_iter(self) -> Self::Iter { + MultiZip { + tuple: ( $( self.$idx.into_par_iter(), )+ ), + } + } + } + + impl<'a, $( $T, )+> IntoParallelIterator for &'a ($( $T, )+) + where + $( + $T: IntoParallelRefIterator<'a>, + $T::Iter: IndexedParallelIterator, + )+ + { + type Item = ($( $T::Item, )+); + type Iter = MultiZip<($( $T::Iter, )+)>; + + fn into_par_iter(self) -> Self::Iter { + MultiZip { + tuple: ( $( self.$idx.par_iter(), )+ ), + } + } + } + + impl<'a, $( $T, )+> IntoParallelIterator for &'a mut ($( $T, )+) + where + $( + $T: IntoParallelRefMutIterator<'a>, + $T::Iter: IndexedParallelIterator, + )+ + { + type Item = ($( $T::Item, )+); + type Iter = MultiZip<($( $T::Iter, )+)>; + + fn into_par_iter(self) -> Self::Iter { + MultiZip { + tuple: ( $( self.$idx.par_iter_mut(), )+ ), + } + } + } + + impl<$( $T, )+> ParallelIterator for MultiZip<($( $T, )+)> + where + $( $T: IndexedParallelIterator, )+ + { + type Item = ($( $T::Item, )+); + + fn drive_unindexed(self, consumer: CONSUMER) -> CONSUMER::Result + where + CONSUMER: UnindexedConsumer, + { + self.drive(consumer) + } + + fn opt_len(&self) -> Option { + Some(self.len()) + } + } + + impl<$( $T, )+> IndexedParallelIterator for MultiZip<($( $T, )+)> + where + $( $T: IndexedParallelIterator, )+ + { + fn drive(self, consumer: CONSUMER) -> CONSUMER::Result + where + CONSUMER: Consumer, + { + zip!($( self.tuple.$idx, )+) + .map(flatten!($( $T, )+)) + .drive(consumer) + } + + fn len(&self) -> usize { + min!($( self.tuple.$idx.len(), )+) + } + + fn with_producer(self, callback: CB) -> CB::Output + where + CB: ProducerCallback, + { + zip!($( self.tuple.$idx, )+) + .map(flatten!($( $T, )+)) + .with_producer(callback) + } + } + )+ + } +} + +multizip_impls! { + Tuple1 { + (0) -> A + } + Tuple2 { + (0) -> A + (1) -> B + } + Tuple3 { + (0) -> A + (1) -> B + (2) -> C + } + Tuple4 { + (0) -> A + (1) -> B + (2) -> C + (3) -> D + } + Tuple5 { + (0) -> A + (1) -> B + (2) -> C + (3) -> D + (4) -> E + } + Tuple6 { + (0) -> A + (1) -> B + (2) -> C + (3) -> D + (4) -> E + (5) -> F + } + Tuple7 { + (0) -> A + (1) -> B + (2) -> C + (3) -> D + (4) -> E + (5) -> F + (6) -> G + } + Tuple8 { + (0) -> A + (1) -> B + (2) -> C + (3) -> D + (4) -> E + (5) -> F + (6) -> G + (7) -> H + } + Tuple9 { + (0) -> A + (1) -> B + (2) -> C + (3) -> D + (4) -> E + (5) -> F + (6) -> G + (7) -> H + (8) -> I + } + Tuple10 { + (0) -> A + (1) -> B + (2) -> C + (3) -> D + (4) -> E + (5) -> F + (6) -> G + (7) -> H + (8) -> I + (9) -> J + } + Tuple11 { + (0) -> A + (1) -> B + (2) -> C + (3) -> D + (4) -> E + (5) -> F + (6) -> G + (7) -> H + (8) -> I + (9) -> J + (10) -> K + } + Tuple12 { + (0) -> A + (1) -> B + (2) -> C + (3) -> D + (4) -> E + (5) -> F + (6) -> G + (7) -> H + (8) -> I + (9) -> J + (10) -> K + (11) -> L + } +} diff --git a/tests/clones.rs b/tests/clones.rs index 1a242b698..f6c1a2237 100644 --- a/tests/clones.rs +++ b/tests/clones.rs @@ -164,3 +164,20 @@ fn clone_repeat() { fn clone_splitter() { check(rayon::iter::split(0..1000, |x| (x, None))); } + +#[test] +fn clone_multizip() { + let v: &Vec<_> = &(0..1000).collect(); + check((v,).into_par_iter()); + check((v, v).into_par_iter()); + check((v, v, v).into_par_iter()); + check((v, v, v, v).into_par_iter()); + check((v, v, v, v, v).into_par_iter()); + check((v, v, v, v, v, v).into_par_iter()); + check((v, v, v, v, v, v, v).into_par_iter()); + check((v, v, v, v, v, v, v, v).into_par_iter()); + check((v, v, v, v, v, v, v, v, v).into_par_iter()); + check((v, v, v, v, v, v, v, v, v, v).into_par_iter()); + check((v, v, v, v, v, v, v, v, v, v, v).into_par_iter()); + check((v, v, v, v, v, v, v, v, v, v, v, v).into_par_iter()); +} diff --git a/tests/debug.rs b/tests/debug.rs index 4a019a8d1..e698d53b7 100644 --- a/tests/debug.rs +++ b/tests/debug.rs @@ -175,3 +175,20 @@ fn debug_repeat() { fn debug_splitter() { check(rayon::iter::split(0..10, |x| (x, None))); } + +#[test] +fn debug_multizip() { + let v: &Vec<_> = &(0..10).collect(); + check((v,).into_par_iter()); + check((v, v).into_par_iter()); + check((v, v, v).into_par_iter()); + check((v, v, v, v).into_par_iter()); + check((v, v, v, v, v).into_par_iter()); + check((v, v, v, v, v, v).into_par_iter()); + check((v, v, v, v, v, v, v).into_par_iter()); + check((v, v, v, v, v, v, v, v).into_par_iter()); + check((v, v, v, v, v, v, v, v, v).into_par_iter()); + check((v, v, v, v, v, v, v, v, v, v).into_par_iter()); + check((v, v, v, v, v, v, v, v, v, v, v).into_par_iter()); + check((v, v, v, v, v, v, v, v, v, v, v, v).into_par_iter()); +} From 963a4d36f33e693f264918ef79402b5db7152284 Mon Sep 17 00:00:00 2001 From: Josh Stone Date: Thu, 5 Dec 2019 11:40:55 -0800 Subject: [PATCH 2/3] Rework the multizip macros for log2 nesting depth --- src/iter/multizip.rs | 55 +++++++++++++++++++++++++++++++++----------- 1 file changed, 41 insertions(+), 14 deletions(-) diff --git a/src/iter/multizip.rs b/src/iter/multizip.rs index bb4f130bb..9284cdfd8 100644 --- a/src/iter/multizip.rs +++ b/src/iter/multizip.rs @@ -82,31 +82,58 @@ pub struct MultiZip { tuple: T, } +// These macros greedily consume 4 or 2 items first to achieve log2 nesting depth. +// For example, 5 => 4,1 => (2,2),1. +// +// The tuples go up to 12, so we might want to greedily consume 8 too, but +// the depth works out the same if we let that expand on the right: +// 9 => 4,5 => (2,2),(4,1) => (2,2),((2,2),1) +// 12 => 4,8 => (2,2),(4,4) => (2,2),((2,2),(2,2)) +// +// But if we ever increase to 13, we would want to split 8,5 rather than 4,9. + macro_rules! zip { - ($first:expr, $( $iter:expr, )*) => { - $first $( .zip($iter) )* + ($a:expr, $b:expr, $c:expr, $d:expr, $( $x:expr, )+) => { + zip!(zip!($a, $b, $c, $d,), zip!($( $x, )+),) + }; + ($a:expr, $b:expr, $( $x:expr, )+) => { + zip!(zip!($a, $b,), zip!($( $x, )+),) + }; + ($a:expr, $( $x:expr, )*) => { + $a $( .zip($x) )* }; } macro_rules! min { - ($x:expr,) => { $x }; - ($x:expr, $( $y:expr, )+) => { cmp::min($x, min!($( $y, )+)) }; + ($a:expr, $b:expr, $c:expr, $d:expr, $( $x:expr, )+) => { + min!(min!($a, $b, $c, $d,), min!($( $x, )+),) + }; + ($a:expr, $b:expr, $( $x:expr, )+) => { + min!(min!($a, $b,), min!($( $x, )+),) + }; + ($a:expr, $b:expr,) => { cmp::min($a, $b) }; + ($a:expr,) => { $a }; +} + +macro_rules! nest { + ($A:tt, $B:tt, $C:tt, $D:tt, $( $X:tt, )+) => { + (nest!($A, $B, $C, $D,), nest!($( $X, )+)) + }; + ($A:tt, $B:tt, $( $X:tt, )+) => { + (($A, $B), nest!($( $X, )+)) + }; + ($A:tt, $B:tt,) => { ($A, $B) }; + ($A:tt,) => { $A }; } macro_rules! flatten { - (|$a:tt : $A:tt| -> ($( $X:ident, )+) { $tuple:tt };) => {{ - fn flatten<$( $X ),+>($a : $A) -> ($( $X, )*) { - $tuple + ($( $T:ident, )+) => {{ + #[allow(non_snake_case)] + fn flatten<$( $T ),+>(nest!($( $T, )+) : nest!($( $T, )+)) -> ($( $T, )+) { + ($( $T, )+) } flatten }}; - (|$a:tt : $A:tt| -> ($( $X:ident, )+) { ($( $x:ident, )+) }; - $B:ident, $( $T:ident, )*) => { - flatten!(|($a, b): ($A, $B)| -> ($( $X, )+ $B,) { ($( $x, )+ b,) }; $( $T, )*) - }; - ($A:ident, $( $T:ident, )*) => { - flatten!(|a: $A| -> ($A,) { (a,) }; $( $T, )*) - }; } macro_rules! multizip_impls { From c0dd154004a14775d5801f85b8372847a87168ba Mon Sep 17 00:00:00 2001 From: Josh Stone Date: Thu, 5 Dec 2019 13:14:33 -0800 Subject: [PATCH 3/3] Simplify multizip reduction macros --- src/iter/multizip.rs | 58 ++++++++++++++++++-------------------------- 1 file changed, 24 insertions(+), 34 deletions(-) diff --git a/src/iter/multizip.rs b/src/iter/multizip.rs index 9284cdfd8..8e36d08a2 100644 --- a/src/iter/multizip.rs +++ b/src/iter/multizip.rs @@ -1,8 +1,6 @@ use super::plumbing::*; use super::*; -use std::cmp; - /// `MultiZip` is an iterator that zips up a tuple of parallel iterators to /// produce tuples of their items. /// @@ -92,44 +90,36 @@ pub struct MultiZip { // // But if we ever increase to 13, we would want to split 8,5 rather than 4,9. -macro_rules! zip { - ($a:expr, $b:expr, $c:expr, $d:expr, $( $x:expr, )+) => { - zip!(zip!($a, $b, $c, $d,), zip!($( $x, )+),) - }; - ($a:expr, $b:expr, $( $x:expr, )+) => { - zip!(zip!($a, $b,), zip!($( $x, )+),) - }; - ($a:expr, $( $x:expr, )*) => { - $a $( .zip($x) )* - }; -} - -macro_rules! min { - ($a:expr, $b:expr, $c:expr, $d:expr, $( $x:expr, )+) => { - min!(min!($a, $b, $c, $d,), min!($( $x, )+),) +macro_rules! reduce { + ($a:expr, $b:expr, $c:expr, $d:expr, $( $x:expr ),+ => $fn:path) => { + reduce!(reduce!($a, $b, $c, $d => $fn), + reduce!($( $x ),+ => $fn) + => $fn) }; - ($a:expr, $b:expr, $( $x:expr, )+) => { - min!(min!($a, $b,), min!($( $x, )+),) + ($a:expr, $b:expr, $( $x:expr ),+ => $fn:path) => { + reduce!(reduce!($a, $b => $fn), + reduce!($( $x ),+ => $fn) + => $fn) }; - ($a:expr, $b:expr,) => { cmp::min($a, $b) }; - ($a:expr,) => { $a }; + ($a:expr, $b:expr => $fn:path) => { $fn($a, $b) }; + ($a:expr => $fn:path) => { $a }; } macro_rules! nest { - ($A:tt, $B:tt, $C:tt, $D:tt, $( $X:tt, )+) => { - (nest!($A, $B, $C, $D,), nest!($( $X, )+)) + ($A:tt, $B:tt, $C:tt, $D:tt, $( $X:tt ),+) => { + (nest!($A, $B, $C, $D), nest!($( $X ),+)) }; - ($A:tt, $B:tt, $( $X:tt, )+) => { - (($A, $B), nest!($( $X, )+)) + ($A:tt, $B:tt, $( $X:tt ),+) => { + (($A, $B), nest!($( $X ),+)) }; - ($A:tt, $B:tt,) => { ($A, $B) }; - ($A:tt,) => { $A }; + ($A:tt, $B:tt) => { ($A, $B) }; + ($A:tt) => { $A }; } macro_rules! flatten { - ($( $T:ident, )+) => {{ + ($( $T:ident ),+) => {{ #[allow(non_snake_case)] - fn flatten<$( $T ),+>(nest!($( $T, )+) : nest!($( $T, )+)) -> ($( $T, )+) { + fn flatten<$( $T ),+>(nest!($( $T ),+) : nest!($( $T ),+)) -> ($( $T, )+) { ($( $T, )+) } flatten @@ -220,21 +210,21 @@ macro_rules! multizip_impls { where CONSUMER: Consumer, { - zip!($( self.tuple.$idx, )+) - .map(flatten!($( $T, )+)) + reduce!($( self.tuple.$idx ),+ => IndexedParallelIterator::zip) + .map(flatten!($( $T ),+)) .drive(consumer) } fn len(&self) -> usize { - min!($( self.tuple.$idx.len(), )+) + reduce!($( self.tuple.$idx.len() ),+ => Ord::min) } fn with_producer(self, callback: CB) -> CB::Output where CB: ProducerCallback, { - zip!($( self.tuple.$idx, )+) - .map(flatten!($( $T, )+)) + reduce!($( self.tuple.$idx ),+ => IndexedParallelIterator::zip) + .map(flatten!($( $T ),+)) .with_producer(callback) } }