Add parFlatTraverseN and parFlatSequenceN

core/shared/src/main/scala/cats/effect/IO.scala

@@ -1599,6 +1599,13 @@ object IO extends IOCompanionPlatform with IOLowPriorityImplicits with TuplePara
   def parTraverseN_[T[_]: Foldable, A, B](n: Int)(ta: T[A])(f: A => IO[B]): IO[Unit] =
     _asyncForIO.parTraverseN_(n)(ta)(f)
 
+  /**
+   * Like `Parallel.parFlatTraverse`, but limits the degree of parallelism.
+   */
+  def parFlatTraverseN[T[_]: Traverse: cats.FlatMap, A, B](n: Int)(ta: T[A])(
+      f: A => IO[T[B]]): IO[T[B]] =
+    _asyncForIO.parFlatTraverseN(n)(ta)(f)
+
   /**
    * Like `Parallel.parSequence`
    */
@@ -1623,6 +1630,12 @@ object IO extends IOCompanionPlatform with IOLowPriorityImplicits with TuplePara
   def parSequenceN_[T[_]: Foldable, A](n: Int)(tma: T[IO[A]]): IO[Unit] =
     _asyncForIO.parSequenceN_(n)(tma)
 
+  /**
+   * Like `Parallel.parFlatSequence`, but limits the degree of parallelism.
+   */
+  def parFlatSequenceN[T[_]: Traverse: cats.FlatMap, A](n: Int)(tmta: T[IO[T[A]]]): IO[T[A]] =
+    _asyncForIO.parFlatSequenceN(n)(tmta)
+
   /**
    * Like `Parallel.parReplicateA`, but limits the degree of parallelism.
    */

kernel/shared/src/main/scala/cats/effect/kernel/GenConcurrent.scala

@@ -16,7 +16,7 @@
 
 package cats.effect.kernel
 
-import cats.{Foldable, Monoid, Semigroup, Traverse}
+import cats.{FlatMap, Foldable, Monoid, Semigroup, Traverse}
 import cats.data.{EitherT, IorT, Kleisli, OptionT, WriterT}
 import cats.effect.kernel.instances.spawn._
 import cats.effect.kernel.syntax.all._
@@ -143,6 +143,18 @@ trait GenConcurrent[F[_], E] extends GenSpawn[F, E] {
    * limit.
    */
   def parTraverseN[T[_]: Traverse, A, B](n: Int)(ta: T[A])(f: A => F[B]): F[T[B]] = {
+    implicit val F: GenConcurrent[F, E] = this
+    parTraverseNImpl[T, A, B, B](n)(ta)(f)(_.sequence[F, B])
+  }
+
+  /**
+   * Shared core implementation for both [[parTraverseN]] and [[parFlatTraverseN]], taking a
+   * function that decides on how to sequence the result.
+   * @tparam B
+   *   the intermediate result type of the function `f`, needs to be sequenceable to `C`
+   */
+  private def parTraverseNImpl[T[_]: Traverse, A, B, C](n: Int)(ta: T[A])(f: A => F[B])(
+      seq: T[F[B]] => F[T[C]]): F[T[C]] = {
     require(n >= 1, s"Concurrency limit should be at least 1, was: $n")
 
     implicit val F: GenConcurrent[F, E] = this
@@ -242,7 +254,7 @@ trait GenConcurrent[F[_], E] extends GenSpawn[F, E] {
               }
             }
 
-            results.flatMap(_.sequence).onCancel(cancelAllAndJoin)
+            results.flatMap(seq).onCancel(cancelAllAndJoin)
           }
       }
     }
@@ -336,6 +348,23 @@ trait GenConcurrent[F[_], E] extends GenSpawn[F, E] {
     }
   }
 
+  /**
+   * Like `Parallel.parFlatSequence`, but limits the degree of parallelism. See [[parSequenceN]]
+   * for fairness considerations, the same semantics apply here.
+   */
+  def parFlatSequenceN[T[_]: Traverse: FlatMap, A](n: Int)(tma: T[F[T[A]]]): F[T[A]] =
+    parFlatTraverseN(n)(tma)(identity)
+
+  /**
+   * Like `Parallel.parFlatTraverse`, but limits the degree of parallelism. See [[parTraverseN]]
+   * for fairness considerations, the same semantics apply here.
+   */
+  def parFlatTraverseN[T[_]: Traverse: FlatMap, A, B](n: Int)(ta: T[A])(
+      f: A => F[T[B]]): F[T[B]] = {
+    implicit val F: GenConcurrent[F, E] = this
+    parTraverseNImpl[T, A, T[B], B](n)(ta)(f)(_.flatSequence)
+  }
+
   override def racePair[A, B](fa: F[A], fb: F[B])
       : F[Either[(Outcome[F, E, A], Fiber[F, E, B]), (Fiber[F, E, A], Outcome[F, E, B])]] = {
     implicit val F: GenConcurrent[F, E] = this

kernel/shared/src/main/scala/cats/effect/kernel/syntax/GenConcurrentSyntax.scala

@@ -16,7 +16,7 @@
 
 package cats.effect.kernel.syntax
 
-import cats.{Foldable, Traverse}
+import cats.{FlatMap, Foldable, Traverse}
 import cats.effect.kernel.GenConcurrent
 
 trait GenConcurrentSyntax {
@@ -34,6 +34,11 @@ trait GenConcurrentSyntax {
   ): ConcurrentParSequenceNOps[T, F, A] =
     new ConcurrentParSequenceNOps(wrapped)
 
+  implicit def concurrentParFlatSequenceOps[T[_], F[_], A](
+      wrapped: T[F[T[A]]]
+  ): ConcurrentParFlatSequenceNOps[T, F, A] =
+    new ConcurrentParFlatSequenceNOps(wrapped)
+
 }
 
 final class GenConcurrentOps_[F[_], A] private[syntax] (private val wrapped: F[A])
@@ -57,6 +62,11 @@ final class ConcurrentParTraverseNOps[T[_], A] private[syntax] (
       f: A => F[B]
   )(implicit T: Foldable[T], F: GenConcurrent[F, ?]): F[Unit] =
     F.parTraverseN_(n)(wrapped)(f)
+
+  def parFlatTraverseN[F[_], B](n: Int)(
+      f: A => F[T[B]]
+  )(implicit T: Traverse[T], FM: FlatMap[T], F: GenConcurrent[F, ?]): F[T[B]] =
+    F.parFlatTraverseN(n)(wrapped)(f)
 }
 
 final class ConcurrentParSequenceNOps[T[_], F[_], A] private[syntax] (
@@ -68,3 +78,11 @@ final class ConcurrentParSequenceNOps[T[_], F[_], A] private[syntax] (
   def parSequenceN_(n: Int)(implicit T: Foldable[T], F: GenConcurrent[F, ?]): F[Unit] =
     F.parSequenceN_(n)(wrapped)
 }
+
+final class ConcurrentParFlatSequenceNOps[T[_], F[_], A] private[syntax] (
+    private val wrapped: T[F[T[A]]]
+) extends AnyVal {
+  def parFlatSequenceN(
+      n: Int)(implicit T: Traverse[T], FM: FlatMap[T], F: GenConcurrent[F, ?]): F[T[A]] =
+    F.parFlatSequenceN(n)(wrapped)
+}

kernel/shared/src/test/scala/cats/effect/kernel/SyntaxSuite.scala

@@ -61,6 +61,16 @@ class SyntaxSuite extends FunSuite {
       result: F[Unit]
     }
 
+    {
+      val result = List(target).parFlatTraverseN(3)(t => t.map(List(_)))
+      result: F[List[A]]
+    }
+
+    {
+      val result = List(target.map(List(_))).parFlatSequenceN(3)
+      result: F[List[A]]
+    }
+
     {
       val result = target.parReplicateAN(3)(5)
       result: F[List[A]]

tests/shared/src/test/scala/cats/effect/IOSuite.scala

@@ -2268,6 +2268,72 @@ class IOSuite extends BaseScalaCheckSuite with DisciplineSuite with IOPlatformSu
     assertCompleteAs(test.attempt.void, ())
   }
 
+  ticked("parFlatTraverseN - run tasks in parallel and flatten the results (ticked)") {
+    implicit ticker =>
+      val p = List(1, 2, 3)
+        .parFlatTraverseN(2) { (n: Int) => IO.sleep(2.seconds) >> IO.pure(List.fill(n)(())) }
+        .timeoutTo(5.seconds, IO(fail("parFlatSequenceN took too long")))
+
+      assertCompleteAs(p, List.fill(1 + 2 + 3)(()))
+  }
+
+  real("parFlatTraverseN - throw when n < 1") {
+    IO.defer {
+      List.empty[Int].parFlatTraverseN(0)(List(_).pure[IO])
+    }.mustFailWith[IllegalArgumentException]
+  }
+
+  real("parFlatTraverseN - propagate errors") {
+    List(1, 2, 3)
+      .parFlatTraverseN(2) { (n: Int) =>
+        if (n == 2) IO.raiseError(new RuntimeException) else List(n).pure[IO]
+      }
+      .mustFailWith[RuntimeException]
+  }
+
+  ticked("parFlatTraverseN - be cancelable") { implicit ticker =>
+    val p = for {
+      f <- List(1, 2, 3).parFlatTraverseN(2)(_ => IO.never[List[Int]]).start
+      _ <- IO.sleep(100.millis)
+      _ <- f.cancel
+    } yield true
+
+    assertCompleteAs(p, true)
+  }
+
+  ticked("parFlatSequenceN - run tasks in parallel and flatten the results (ticked)") {
+    implicit ticker =>
+      val p = List(1, 2, 3)
+        .map { (n: Int) => IO.sleep(2.seconds) >> IO.pure(List.fill(n)(())) }
+        .parFlatSequenceN(2)
+        .timeoutTo(5.seconds, IO(fail("parFlatSequenceN took too long")))
+
+      assertCompleteAs(p, List.fill(1 + 2 + 3)(()))
+  }
+
+  real("parFlatSequenceN - throw when n < 1") {
+    IO.defer {
+      List.empty[IO[List[Int]]].parFlatSequenceN(0)
+    }.mustFailWith[IllegalArgumentException]
+  }
+
+  real("parFlatSequenceN - propagate errors") {
+    List(1, 2, 3)
+      .map { (n: Int) => if (n == 2) IO.raiseError(new RuntimeException) else List(n).pure[IO] }
+      .parFlatSequenceN(2)
+      .mustFailWith[RuntimeException]
+  }
+
+  ticked("parFlatSequenceN - be cancelable") { implicit ticker =>
+    val p = for {
+      f <- List(1, 2, 3).map(_ => IO.never[List[IO[Int]]]).parFlatSequenceN(2).start
+      _ <- IO.sleep(100.millis)
+      _ <- f.cancel
+    } yield true
+
+    assertCompleteAs(p, true)
+  }
+
   real("parallel - run parallel actually in parallel") {
     val x = IO.sleep(2.seconds) >> IO.pure(1)
     val y = IO.sleep(2.seconds) >> IO.pure(2)