Merge pull request #4243 from TonioGela/main

Added `parFlatMapN`
typelevel · Oct 7, 2022 · aadce06 · aadce06
2 parents 54bbe14 + 325b4b4
commit aadce06
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 4 deletions.
diff --git a/project/Boilerplate.scala b/project/Boilerplate.scala
@@ -341,6 +341,9 @@ object Boilerplate {
         -  /** @group ParMapArity */
         -  def parMap$arity[M[_], ${`A..N`}, Z]($fparams)(f: (${`A..N`}) => Z)(implicit p: NonEmptyParallel[M]): M[Z] =
         -    p.flatMap.map(${nestedExpansion.products}) { case ${nestedExpansion.`(a..n)`} => f(${`a..n`}) }
+        -
+        -  def parFlatMap$arity[M[_], ${`A..N`}, Z]($fparams)(f: (${`A..N`}) => M[Z])(implicit p: NonEmptyParallel[M]): M[Z] =
+        -    p.flatMap.flatMap(${nestedExpansion.products}) { case ${nestedExpansion.`(a..n)`} => f(${`a..n`}) }
       |}
       """
     }
@@ -464,6 +467,11 @@ object Boilerplate {
         else
           s"def parMapN[Z](f: (${`A..N`}) => Z)(implicit p: NonEmptyParallel[M]): M[Z] = Parallel.parMap$arity($tupleArgs)(f)"
 
+      val parFlatMap =
+        if (arity == 1)
+          s"def parFlatMap[Z](f: (${`A..N`}) => M[Z])(implicit p: NonEmptyParallel[M]): M[Z] = p.flatMap.flatMap($tupleArgs)(f)"
+        else
+          s"def parFlatMapN[Z](f: (${`A..N`}) => M[Z])(implicit p: NonEmptyParallel[M]): M[Z] = Parallel.parFlatMap$arity($tupleArgs)(f)"
       val parTupled =
         if (arity == 1) ""
         else
@@ -482,6 +490,7 @@ object Boilerplate {
          -private[syntax] final class Tuple${arity}ParallelOps[M[_], ${`A..N`}](private val $tupleTpe) extends Serializable {
          -  $parMap
          -  $parTupled
+         -  $parFlatMap
          -}
       |
       """

diff --git a/tests/shared/src/test/scala/cats/tests/ParallelSuite.scala b/tests/shared/src/test/scala/cats/tests/ParallelSuite.scala
@@ -390,6 +390,38 @@ class ParallelSuite
     }
   }
 
+  test("ParFlatMapN over List should be consistent with parMapN flatten") {
+    forAll { (as: List[Int], bs: List[Int], cs: List[Int], mf: (Int, Int, Int) => List[Int]) =>
+      assert((as, bs, cs).parFlatMapN(mf) == (as, bs, cs).parMapN(mf).flatten)
+    }
+  }
+
+  test("ParFlatMapN over NonEmptyList should be consistent with parMapN flatten") {
+    forAll {
+      (as: NonEmptyList[Int], bs: NonEmptyList[Int], cs: NonEmptyList[Int], mf: (Int, Int, Int) => NonEmptyList[Int]) =>
+        assert((as, bs, cs).parFlatMapN(mf) == (as, bs, cs).parMapN(mf).flatten)
+    }
+  }
+
+  test("ParFlatMap over List should be consistent with flatmap") {
+    forAll { (as: List[Int], mf: Int => List[Int]) =>
+      assert(Tuple1(as).parFlatMap(mf) == Tuple1(as).flatMap(mf))
+    }
+  }
+
+  test("ParFlatMap over NonEmptyList should be consistent with flatmap") {
+    forAll { (as: NonEmptyList[Int], mf: Int => NonEmptyList[Int]) =>
+      assert(Tuple1(as).parFlatMap(mf) == Tuple1(as).flatMap(mf))
+    }
+  }
+
+  test("ParMapN over f should be consistent with parFlatMapN over f lifted in List") {
+    forAll { (as: List[Int], bs: List[Int], f: (Int, Int) => Int) =>
+      val mf: (Int, Int) => List[Int] = (a, b) => f(a, b).pure[List]
+      assert((as, bs).parMapN(f) == (as, bs).parFlatMapN(mf))
+    }
+  }
+
   test("ParTupled of NonEmptyList should be consistent with ParMap of Tuple.apply") {
     forAll { (fa: NonEmptyList[Int], fb: NonEmptyList[Int], fc: NonEmptyList[Int], fd: NonEmptyList[Int]) =>
       assert((fa, fb, fc, fd).parTupled === ((fa, fb, fc, fd).parMapN(Tuple4.apply)))

diff --git a/tests/shared/src/test/scala/cats/tests/SyntaxSuite.scala b/tests/shared/src/test/scala/cats/tests/SyntaxSuite.scala
@@ -241,14 +241,32 @@ object SyntaxSuite {
   }
 
   def testParallelTuple[M[_]: Monad, F[_], A, B, C, Z](implicit P: NonEmptyParallel.Aux[M, F]) = {
-    val tfabc = mock[(M[A], M[B], M[C])]
     val fa = mock[M[A]]
     val fb = mock[M[B]]
     val fc = mock[M[C]]
-    val f = mock[(A, B, C) => Z]
 
-    tfabc.parMapN(f)
-    (fa, fb, fc).parMapN(f)
+    val tfabc = mock[(M[A], M[B], M[C])]
+    val fthree = mock[(A, B, C) => Z]
+    val mfthree = mock[(A, B, C) => M[Z]]
+
+    tfabc.parMapN(fthree)
+    (fa, fb, fc).parMapN(fthree)
+    tfabc.parFlatMapN(mfthree)
+    (fa, fb, fc).parFlatMapN(mfthree)
+
+    val tfab = mock[(M[A], M[B])]
+    val ftwo = mock[(A, B) => Z]
+    val mftwo = mock[(A, B) => M[Z]]
+
+    tfab.parMapN(ftwo)
+    (fa, fb).parMapN(ftwo)
+    tfab.parFlatMapN(mftwo)
+    (fa, fb).parFlatMapN(mftwo)
+
+    val tfa = mock[Tuple1[M[A]]]
+    val mfone = mock[A => M[Z]]
+
+    tfa.parFlatMap(mfone)
   }
 
   def testParallelBi[M[_], F[_], T[_, _]: Bitraverse, A, B, C, D](implicit P: Parallel.Aux[M, F]): Unit = {