diff --git a/src/iter/mod.rs b/src/iter/mod.rs index 36294cf12..7d61022ef 100644 --- a/src/iter/mod.rs +++ b/src/iter/mod.rs @@ -1519,6 +1519,20 @@ pub trait ParallelIterator: Sized + Send { /// assert_eq!(left, [0, 1, 2, 3]); /// assert_eq!(right, [1, 2, 3, 4]); /// ``` + /// + /// Nested pairs can be unzipped too. + /// + /// ``` + /// use rayon::prelude::*; + /// + /// let (values, (squares, cubes)): (Vec<_>, (Vec<_>, Vec<_>)) = (0..4).into_par_iter() + /// .map(|i| (i, (i * i, i * i * i))) + /// .unzip(); + /// + /// assert_eq!(values, [0, 1, 2, 3]); + /// assert_eq!(squares, [0, 1, 4, 9]); + /// assert_eq!(cubes, [0, 1, 8, 27]); + /// ``` fn unzip(self) -> (FromA, FromB) where Self: ParallelIterator, FromA: Default + Send + ParallelExtend, @@ -1567,17 +1581,38 @@ pub trait ParallelIterator: Sized + Send { /// use rayon::iter::Either; /// /// let (left, right): (Vec<_>, Vec<_>) = (0..8).into_par_iter() - /// .partition_map(|x| { - /// if x % 2 == 0 { - /// Either::Left(x * 4) - /// } else { - /// Either::Right(x * 3) - /// } - /// }); + /// .partition_map(|x| { + /// if x % 2 == 0 { + /// Either::Left(x * 4) + /// } else { + /// Either::Right(x * 3) + /// } + /// }); /// /// assert_eq!(left, [0, 8, 16, 24]); /// assert_eq!(right, [3, 9, 15, 21]); /// ``` + /// + /// Nested `Either` enums can be split as well. + /// + /// ``` + /// use rayon::prelude::*; + /// use rayon::iter::Either::*; + /// + /// let ((fizzbuzz, fizz), (buzz, other)): ((Vec<_>, Vec<_>), (Vec<_>, Vec<_>)) = (1..20) + /// .into_par_iter() + /// .partition_map(|x| match (x % 3, x % 5) { + /// (0, 0) => Left(Left(x)), + /// (0, _) => Left(Right(x)), + /// (_, 0) => Right(Left(x)), + /// (_, _) => Right(Right(x)), + /// }); + /// + /// assert_eq!(fizzbuzz, [15]); + /// assert_eq!(fizz, [3, 6, 9, 12, 18]); + /// assert_eq!(buzz, [5, 10]); + /// assert_eq!(other, [1, 2, 4, 7, 8, 11, 13, 14, 16, 17, 19]); + /// ``` fn partition_map(self, predicate: P) -> (A, B) where A: Default + Send + ParallelExtend, B: Default + Send + ParallelExtend, diff --git a/src/iter/unzip.rs b/src/iter/unzip.rs index 9f10bd713..3bf75bc97 100644 --- a/src/iter/unzip.rs +++ b/src/iter/unzip.rs @@ -24,7 +24,7 @@ trait UnzipOp: Sync + Send { } } -/// Run an unzip-like operation into `ParallelExtend` collections. +/// Run an unzip-like operation into default `ParallelExtend` collections. fn execute(pi: I, op: OP) -> (FromA, FromB) where I: ParallelIterator, OP: UnzipOp, @@ -33,21 +33,30 @@ fn execute(pi: I, op: OP) -> (FromA, FromB) { let mut a = FromA::default(); let mut b = FromB::default(); - { - // We have no idea what the consumers will look like for these - // collections' `par_extend`, but we can intercept them in our own - // `drive_unindexed`. Start with the left side, type `A`: - let iter = UnzipA { - base: pi, - op: op, - b: &mut b, - }; - a.par_extend(iter); - } + execute_into(&mut a, &mut b, pi, op); (a, b) } +/// Run an unzip-like operation into `ParallelExtend` collections. +fn execute_into(a: &mut FromA, b: &mut FromB, pi: I, op: OP) + where I: ParallelIterator, + OP: UnzipOp, + FromA: Send + ParallelExtend, + FromB: Send + ParallelExtend +{ + // We have no idea what the consumers will look like for these + // collections' `par_extend`, but we can intercept them in our own + // `drive_unindexed`. Start with the left side, type `A`: + let iter = UnzipA { + base: pi, + op: op, + b: b, + }; + a.par_extend(iter); +} + + /// Unzips the items of a parallel iterator into a pair of arbitrary /// `ParallelExtend` containers. /// @@ -188,7 +197,7 @@ struct UnzipA<'b, I, OP, FromB: 'b> { impl<'b, I, OP, FromB> ParallelIterator for UnzipA<'b, I, OP, FromB> where I: ParallelIterator, OP: UnzipOp, - FromB: Default + Send + ParallelExtend + FromB: Send + ParallelExtend { type Item = OP::Left; @@ -386,3 +395,57 @@ impl Reducer<(A, B)> for UnzipReducer (self.left.reduce(left.0, right.0), self.right.reduce(left.1, right.1)) } } + + +impl ParallelExtend<(A, B)> for (FromA, FromB) +where + A: Send, + B: Send, + FromA: Send + ParallelExtend, + FromB: Send + ParallelExtend, +{ + fn par_extend(&mut self, pi: I) + where + I: IntoParallelIterator, + { + execute_into(&mut self.0, &mut self.1, pi.into_par_iter(), Unzip); + } +} + +impl ParallelExtend> for (A, B) +where + L: Send, + R: Send, + A: Send + ParallelExtend, + B: Send + ParallelExtend, +{ + fn par_extend(&mut self, pi: I) + where + I: IntoParallelIterator>, + { + execute_into(&mut self.0, &mut self.1, pi.into_par_iter(), UnEither); + } +} + +/// An `UnzipOp` that routes items depending on their `Either` variant. +struct UnEither; + +impl UnzipOp> for UnEither +where + L: Send, + R: Send, +{ + type Left = L; + type Right = R; + + fn consume(&self, item: Either, left: FL, right: FR) -> (FL, FR) + where + FL: Folder, + FR: Folder, + { + match item { + Either::Left(item) => (left.consume(item), right), + Either::Right(item) => (left, right.consume(item)), + } + } +}