Add cartesian_product to ParallelIterator

src/iter/cartesian_product.rs

@@ -0,0 +1,48 @@
+use super::plumbing::*;
+use super::repeat::*;
+use super::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
+
+/// `CartesianProduct` is an iterator that combines `i` and `j` into a single iterator that
+/// iterates over the cartesian product of the elements of `i` and `j`. This struct is created by
+/// the [`cartesian_product()`] method on [`ParallelIterator`]
+#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
+#[derive(Debug, Clone, Copy)]
+pub struct CartesianProduct<I, J>
+where
+    I: ParallelIterator,
+    J: ParallelIterator,
+{
+    i: I,
+    j: J,
+}
+
+impl<I, J> CartesianProduct<I, J>
+where
+    I: ParallelIterator,
+    J: IndexedParallelIterator,
+{
+    /// Creates a new `CartesianProduct` iterator.
+    pub(super) fn new(i: I, j: J) -> Self {
+        CartesianProduct { i, j }
+    }
+}
+
+impl<I, J> ParallelIterator for CartesianProduct<I, J>
+where
+    I: ParallelIterator,
+    J: IndexedParallelIterator + Clone + Sync,
+    I::Item: Clone + Send,
+{
+    type Item = (I::Item, J::Item);
+
+    fn drive_unindexed<C>(self, consumer: C) -> C::Result
+    where
+        C: UnindexedConsumer<Self::Item>,
+    {
+        self.i
+            .into_par_iter()
+            .map(|i_item| repeat(i_item.clone()).zip(self.j.clone()))
+            .flatten()
+            .drive_unindexed(consumer)
+    }
+}

src/iter/mod.rs

@@ -104,6 +104,7 @@ mod test;
 //   can be readily distinguished.
 
 mod blocks;
+mod cartesian_product;
 mod chain;
 mod chunks;
 mod cloned;
@@ -163,6 +164,7 @@ mod zip_eq;
 
 pub use self::{
     blocks::{ExponentialBlocks, UniformBlocks},
+    cartesian_product::CartesianProduct,
     chain::Chain,
     chunks::Chunks,
     cloned::Cloned,
@@ -606,6 +608,31 @@ pub trait ParallelIterator: Sized + Send {
         Map::new(self, map_op)
     }
 
+    /// Produces a new iterator that iterates over the cartesian product of the element sets of
+    /// `self` and `CI`.
+    ///
+    /// Iterator element type is `(Self::Item, CI::Item)`
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use rayon::prelude::*;
+    ///
+    /// let set: Vec<_> = (0..2).into_par_iter().cartesian_product(0..2).collect();
+    /// assert!(set.contains(&(0, 0)));
+    /// assert!(set.contains(&(0, 1)));
+    /// assert!(set.contains(&(1, 0)));
+    /// assert!(set.contains(&(1, 1)));
+    /// ```
+    fn cartesian_product<CI>(self, other: CI) -> CartesianProduct<Self, CI::Iter>
+    where
+        Self::Item: Clone + Send,
+        CI: IntoParallelIterator,
+        CI::Iter: IndexedParallelIterator + Clone + Sync,
+    {
+        CartesianProduct::new(self, other.into_par_iter())
+    }
+
     /// Applies `map_op` to the given `init` value with each item of this
     /// iterator, producing a new iterator with the results.
     ///

tests/cartesian_product.rs

@@ -0,0 +1,53 @@
+use rayon::prelude::*;
+use std::collections::HashSet;
+
+#[test]
+fn cartesian_ranges() {
+    let a: Vec<(usize, usize)> = (0..3).into_par_iter().cartesian_product(0..3).collect();
+    let b: HashSet<(usize, usize)> = HashSet::from_iter(a.into_iter());
+    assert!(b.contains(&(0, 0)));
+    assert!(b.contains(&(0, 1)));
+    assert!(b.contains(&(0, 2)));
+    assert!(b.contains(&(1, 0)));
+    assert!(b.contains(&(1, 1)));
+    assert!(b.contains(&(1, 2)));
+    assert!(b.contains(&(2, 0)));
+    assert!(b.contains(&(2, 1)));
+    assert!(b.contains(&(2, 2)));
+}
+
+#[test]
+fn cartesian_vecs() {
+    let a: Vec<(f64, f64)> = vec![0.1, 1.2, 2.4]
+        .into_par_iter()
+        .cartesian_product(vec![4.8, 16.32, 32.64])
+        .collect();
+    assert!(a.contains(&(0.1, 4.8)));
+    assert!(a.contains(&(0.1, 16.32)));
+    assert!(a.contains(&(0.1, 32.64)));
+    assert!(a.contains(&(1.2, 4.8)));
+    assert!(a.contains(&(1.2, 16.32)));
+    assert!(a.contains(&(1.2, 32.64)));
+    assert!(a.contains(&(2.4, 4.8)));
+    assert!(a.contains(&(2.4, 16.32)));
+    assert!(a.contains(&(2.4, 32.64)));
+}
+
+#[test]
+fn cartesian_chain() {
+    let a: Vec<(usize, usize, usize)> = (0..2)
+        .into_par_iter()
+        .cartesian_product(0..2)
+        .cartesian_product(0..2)
+        .map(|((a, b), c)| (a, b, c))
+        .collect();
+    let b: HashSet<(usize, usize, usize)> = HashSet::from_iter(a.into_iter());
+    assert!(b.contains(&(0, 0, 0)));
+    assert!(b.contains(&(0, 0, 1)));
+    assert!(b.contains(&(0, 1, 0)));
+    assert!(b.contains(&(0, 1, 1)));
+    assert!(b.contains(&(1, 0, 0)));
+    assert!(b.contains(&(1, 0, 1)));
+    assert!(b.contains(&(1, 1, 0)));
+    assert!(b.contains(&(1, 1, 1)));
+}