Update Signature of ForEach#collectM (#1236)

* update signature of foreach collectm

* format

benchmarks/src/main/scala/zio/prelude/CollectBenchmark.scala

@@ -0,0 +1,53 @@
+package zio.prelude
+
+import cats.effect.unsafe.implicits.global
+import cats.effect.{IO => CIO}
+import cats.instances.list._
+import cats.syntax.all._
+import org.openjdk.jmh.annotations.{State => BenchmarkState, _}
+import org.openjdk.jmh.infra.Blackhole
+import zio.prelude.fx.ZPure
+import zio.{Scope => _, _}
+
+import java.util.concurrent.TimeUnit
+
+@BenchmarkState(Scope.Thread)
+@BenchmarkMode(Array(Mode.Throughput))
+@OutputTimeUnit(TimeUnit.SECONDS)
+@Measurement(iterations = 5, timeUnit = TimeUnit.SECONDS, time = 3)
+@Warmup(iterations = 5, timeUnit = TimeUnit.SECONDS, time = 3)
+@Fork(value = 1)
+class CollectBenchmarks {
+
+  var list: List[Int] = _
+
+  @Param(Array("100000"))
+  var size: Int = _
+
+  @Setup(Level.Trial)
+  def setup(): Unit =
+    list = (1 to size).toList
+
+  @Benchmark
+  def catsTraverseFilterOption(bh: Blackhole): Unit =
+    bh.consume(list.traverse(n => Option(Option(n))))
+
+  @Benchmark
+  def catsTraverseFilterCIO(bh: Blackhole): Unit =
+    bh.consume(list.traverseFilter(n => CIO.pure(Option(n))).unsafeRunSync())
+
+  @Benchmark
+  def zioCollectMOption(bh: Blackhole): Unit =
+    bh.consume(list.collectM(n => Option(Option(n))))
+
+  @Benchmark
+  def zioCollectMZIO(bh: Blackhole): Unit =
+    bh.consume(unsafeRun(list.collectM(n => ZIO.succeed(Option(n)))))
+
+  @Benchmark
+  def zioCollectMZPure(bh: Blackhole): Unit =
+    bh.consume(list.collectM(n => ZPure.succeed[Unit, Option[Int]](Option(n))).run)
+
+  def unsafeRun[E, A](zio: ZIO[Any, E, A]): A =
+    Unsafe.unsafe(implicit unsafe => Runtime.default.unsafe.run(zio).getOrThrowFiberFailure())
+}

core-tests/shared/src/test/scala/zio/prelude/ForEachSpec.scala

@@ -39,6 +39,9 @@ object ForEachSpec extends ZIOBaseSpec {
   val genEitherIntIntFunction: Gen[Any, Int => Either[Int, Int]] =
     Gen.function(Gen.either(genInt, genInt))
 
+  val genIntPartialFunction: Gen[Any, PartialFunction[Int, Int]] =
+    Gen.partialFunction(Gen.int)
+
   implicit val chunkOptionForEach: ForEach[ChunkOption] =
     ForEach[Chunk].compose[Option]
 
@@ -54,6 +57,13 @@ object ForEachSpec extends ZIOBaseSpec {
         test("vector")(checkAllLaws(ForEachLaws)(GenF.vector, Gen.int))
       ),
       suite("combinators")(
+        test("collect") {
+          check(genList, genIntPartialFunction) { (as, pf) =>
+            val actual   = ForEach[List].collect(as)(pf)
+            val expected = as.collect(pf)
+            assert(actual)(equalTo(expected))
+          }
+        },
         test("contains") {
           check(genList, genInt) { (as, a) =>
             val actual   = ForEach[List].contains(as)(a)

core/shared/src/main/scala-2.12/zio/prelude/InvariantVersionSpecific.scala

@@ -29,6 +29,10 @@ trait InvariantVersionSpecific {
     new ForEach[F] {
       def forEach[G[+_]: IdentityBoth: Covariant, A, B](fa: F[A])(f: A => G[B]): G[F[B]]          =
         CovariantIdentityBoth[G].forEach(fa)(f)(derive.derive)
+      override def collectM[G[+_]: IdentityBoth: Covariant, A, B](fa: F[A])(
+        f: A => G[Option[B]]
+      )(implicit identityBoth: IdentityBoth[F], identityEither: IdentityEither[F]): G[F[B]] =
+        CovariantIdentityBoth[G].collectM(fa)(f)(derive.derive)
       override def forEach_[G[+_]: IdentityBoth: Covariant, A](fa: F[A])(f: A => G[Any]): G[Unit] =
         CovariantIdentityBoth[G].forEach_(fa)(f)
     }

core/shared/src/main/scala-2.13+/zio/prelude/InvariantVersionSpecific.scala

@@ -29,6 +29,10 @@ trait InvariantVersionSpecific {
     new ForEach[F] {
       def forEach[G[+_]: IdentityBoth: Covariant, A, B](fa: F[A])(f: A => G[B]): G[F[B]]          =
         CovariantIdentityBoth[G].forEach(fa)(f)(derive.derive)
+      override def collectM[G[+_]: IdentityBoth: Covariant, A, B](fa: F[A])(
+        f: A => G[Option[B]]
+      )(implicit identityBoth: IdentityBoth[F], identityEither: IdentityEither[F]): G[F[B]] =
+        CovariantIdentityBoth[G].collectM(fa)(f)(derive.derive)
       override def forEach_[G[+_]: IdentityBoth: Covariant, A](fa: F[A])(f: A => G[Any]): G[Unit] =
         CovariantIdentityBoth[G].forEach_(fa)(f)
     }

core/shared/src/main/scala/zio/prelude/AssociativeBoth.scala

@@ -1234,6 +1234,20 @@ object AssociativeBoth extends AssociativeBothLowPriority {
         fa.zipWithPar(fb)((_, _))
       def map[A, B](f: A => B): ZIO[R, E, A] => ZIO[R, E, B]                                =
         _.map(f)
+      override def collectM[A, B, Collection[+Element] <: Iterable[Element]](in: Collection[A])(
+        f: A => ZIO[R, E, Option[B]]
+      )(implicit bf: BuildFrom[Collection[A], B, Collection[B]]): ZIO[R, E, Collection[B]] =
+        ZIO.suspendSucceed {
+          val iterator = in.iterator
+          val builder  = bf.newBuilder(in)
+
+          ZIO
+            .whileLoop(iterator.hasNext)(f(iterator.next())) {
+              case Some(b) => builder += b
+              case None    =>
+            }
+            .as(builder.result())
+        }
       override def forEach[A, B, Collection[+Element] <: Iterable[Element]](in: Collection[A])(f: A => ZIO[R, E, B])(
         implicit bf: BuildFrom[Collection[A], B, Collection[B]]
       ): ZIO[R, E, Collection[B]] =

core/shared/src/main/scala/zio/prelude/ForEach.scala

@@ -44,30 +44,22 @@ trait ForEach[F[+_]] extends Covariant[F] { self =>
    */
   def collect[A, B](
     fa: F[A]
-  )(pf: PartialFunction[A, B])(implicit identityBoth: IdentityBoth[F], identityEither: IdentityEither[F]): F[B] = {
-    implicit val covariant: Covariant[F]  = self
-    implicit val identity: Identity[F[B]] = Identity.fromIdentityEitherCovariant
-    foldMap(fa) { a =>
-      pf.lift(a) match {
-        case Some(b) => b.succeed[F]
-        case None    => identityEither.none
-      }
-    }
-  }
+  )(pf: PartialFunction[A, B])(implicit identityBoth: IdentityBoth[F], identityEither: IdentityEither[F]): F[B] =
+    Id.unwrap(collectM(fa)(pf.lift.andThen(Id(_))))
 
   /**
    * Collects elements of the collection for which the effectual partial function
    * `pf` is defined.
    */
-  def collectM[G[+_]: IdentityFlatten: Covariant, A, B](fa: F[A])(
-    pf: PartialFunction[A, G[B]]
+  def collectM[G[+_]: IdentityBoth: Covariant, A, B](fa: F[A])(
+    f: A => G[Option[B]]
   )(implicit identityBoth: IdentityBoth[F], identityEither: IdentityEither[F]): G[F[B]] = {
     implicit val covariant: Covariant[F]  = self
     implicit val identity: Identity[F[B]] = Identity.fromIdentityEitherCovariant
     foldMapM(fa) { a =>
-      pf.lift(a) match {
-        case Some(fb) => fb.map(_.succeed[F])
-        case None     => IdentityFlatten[G].any.as(identityEither.none)
+      f(a).map {
+        case Some(b) => b.succeed[F]
+        case None    => identityEither.none
       }
     }
   }
@@ -104,28 +96,16 @@ trait ForEach[F[+_]] extends Covariant[F] { self =>
    */
   def filter[A](
     fa: F[A]
-  )(f: A => Boolean)(implicit identityBoth: IdentityBoth[F], identityEither: IdentityEither[F]): F[A] = {
-    implicit val covariant: Covariant[F]  = self
-    implicit val identity: Identity[F[A]] = Identity.fromIdentityEitherCovariant
-    foldMap(fa) { a =>
-      if (f(a)) a.succeed[F] else identityEither.none
-    }
-  }
+  )(f: A => Boolean)(implicit identityBoth: IdentityBoth[F], identityEither: IdentityEither[F]): F[A] =
+    Id.unwrap(filterM(fa)(a => Id(f(a))))
 
   /**
    * Filters the collection with the effectual predicate `f`.
    */
-  def filterM[G[+_]: IdentityFlatten: Covariant, A](
+  def filterM[G[+_]: IdentityBoth: Covariant, A](
     fa: F[A]
-  )(f: A => G[Boolean])(implicit identityBoth: IdentityBoth[F], identityEither: IdentityEither[F]): G[F[A]] = {
-    implicit val covariant: Covariant[F]  = self
-    implicit val identity: Identity[F[A]] = Identity.fromIdentityEitherCovariant
-    foldMapM(fa) { a =>
-      f(a).map { b =>
-        if (b) a.succeed[F] else identityEither.none
-      }
-    }
-  }
+  )(f: A => G[Boolean])(implicit identityBoth: IdentityBoth[F], identityEither: IdentityEither[F]): G[F[A]] =
+    collectM(fa)(a => f(a).map(b => if (b) Some(a) else None))
 
   /**
    * Returns the first element in the collection satisfying the specified
@@ -179,8 +159,8 @@ trait ForEach[F[+_]] extends Covariant[F] { self =>
    * a single summary using the `combine` operation of `Identity`, or the
    * `identity` element if the collection is empty.
    */
-  def foldMapM[G[+_]: Covariant: IdentityFlatten, A, B: Identity](fa: F[A])(f: A => G[B]): G[B] =
-    foldLeftM[G, B, A](fa)(Identity[B].identity)((accu, a) => f(a).map(aa => accu.combine(aa)))
+  def foldMapM[G[+_]: Covariant: IdentityBoth, A, B: Identity](fa: F[A])(f: A => G[B]): G[B] =
+    forEach(fa)(f).map(fold[B])
 
   /**
    * Folds over the elements of this collection from right to left to produce a
@@ -336,16 +316,16 @@ trait ForEach[F[+_]] extends Covariant[F] { self =>
   /**
    * Partitions the collection based on the specified effectual function.
    */
-  def partitionMapM[G[+_]: IdentityFlatten: Covariant, A, B, C](
+  def partitionMapM[G[+_]: IdentityBoth: Covariant, A, B, C](
     fa: F[A]
   )(f: A => G[Either[B, C]])(implicit both: IdentityBoth[F], either: IdentityEither[F]): G[(F[B], F[C])] = {
     implicit val covariant: Covariant[F]       = self
     implicit val leftIdentity: Identity[F[B]]  = Identity.fromIdentityEitherCovariant
     implicit val rightIdentity: Identity[F[C]] = Identity.fromIdentityEitherCovariant
     foldMapM(fa) { a =>
       f(a).map {
-        case Left(b)  => (b.succeed, either.none)
-        case Right(c) => (either.none, c.succeed)
+        case Left(b)  => (b.succeed[F], either.none)
+        case Right(c) => (either.none, c.succeed[F])
       }
     }
   }
@@ -488,12 +468,12 @@ trait ForEachSyntax {
       B: IdentityBoth[F]
     ): F[B] =
       F.collect(self)(pf)
-    def collectM[G[+_]: IdentityFlatten: Covariant, B](pf: PartialFunction[A, G[B]])(implicit
+    def collectM[G[+_]: IdentityBoth: Covariant, B](f: A => G[Option[B]])(implicit
       F: ForEach[F],
       I: IdentityEither[F],
       B: IdentityBoth[F]
     ): G[F[B]] =
-      F.collectM(self)(pf)
+      F.collectM(self)(f)
     def concatenate(implicit F: ForEach[F], A: Identity[A]): A                                                  =
       F.concatenate(self)
     def forEach[G[+_]: IdentityBoth: Covariant, B](f: A => G[B])(implicit F: ForEach[F]): G[F[B]]               =
@@ -506,7 +486,7 @@ trait ForEachSyntax {
       F.exists(self)(f)
     def filter(f: A => Boolean)(implicit F: ForEach[F], I: IdentityEither[F], B: IdentityBoth[F]): F[A]         =
       F.filter(self)(f)
-    def filterM[G[+_]: IdentityFlatten: Covariant](f: A => G[Boolean])(implicit
+    def filterM[G[+_]: IdentityBoth: Covariant](f: A => G[Boolean])(implicit
       F: ForEach[F],
       I: IdentityEither[F],
       B: IdentityBoth[F]
@@ -520,7 +500,7 @@ trait ForEachSyntax {
       F.foldLeftM(self)(s)(f)
     def foldMap[B: Identity](f: A => B)(implicit F: ForEach[F]): B                                              =
       F.foldMap(self)(f)
-    def foldMapM[G[+_]: Covariant: IdentityFlatten, B: Identity](f: A => G[B])(implicit F: ForEach[F]): G[B]    =
+    def foldMapM[G[+_]: Covariant: IdentityBoth, B: Identity](f: A => G[B])(implicit F: ForEach[F]): G[B]       =
       F.foldMapM(self)(f)
     def foldRight[S](s: S)(f: (A, S) => S)(implicit F: ForEach[F]): S                                           =
       F.foldRight(self)(s)(f)

core/shared/src/main/scala/zio/prelude/Invariant.scala

@@ -78,6 +78,10 @@ object Invariant extends LowPriorityInvariantImplicits with InvariantVersionSpec
     new ForEach[Chunk] {
       def forEach[G[+_]: IdentityBoth: Covariant, A, B](chunk: Chunk[A])(f: A => G[B]): G[Chunk[B]]      =
         CovariantIdentityBoth[G].forEach(chunk)(f)
+      override def collectM[G[+_]: IdentityBoth: Covariant, A, B](chunk: Chunk[A])(
+        f: A => G[Option[B]]
+      )(implicit identityBoth: IdentityBoth[Chunk], identityEither: IdentityEither[Chunk]): G[Chunk[B]] =
+        CovariantIdentityBoth[G].collectM(chunk)(f)
       override def forEach_[G[+_]: IdentityBoth: Covariant, A](chunk: Chunk[A])(f: A => G[Any]): G[Unit] =
         CovariantIdentityBoth[G].forEach_(chunk)(f)
     }
@@ -679,6 +683,10 @@ object Invariant extends LowPriorityInvariantImplicits with InvariantVersionSpec
     new ForEach[List] {
       def forEach[G[+_]: IdentityBoth: Covariant, A, B](list: List[A])(f: A => G[B]): G[List[B]]     =
         CovariantIdentityBoth[G].forEach(list)(f)
+      override def collectM[G[+_]: IdentityBoth: Covariant, A, B](fa: List[A])(
+        f: A => G[Option[B]]
+      )(implicit identityBoth: IdentityBoth[List], identityEither: IdentityEither[List]): G[List[B]] =
+        CovariantIdentityBoth[G].collectM(fa)(f)
       override def forEach_[G[+_]: IdentityBoth: Covariant, A](fa: List[A])(f: A => G[Any]): G[Unit] =
         CovariantIdentityBoth[G].forEach_(fa)(f)
       override def map[A, B](f: A => B): List[A] => List[B]                                          =
@@ -694,6 +702,15 @@ object Invariant extends LowPriorityInvariantImplicits with InvariantVersionSpec
         CovariantIdentityBoth[G]
           .forEach[(K, V), (K, V2), Iterable](map) { case (k, v) => f(v).map(k -> _) }
           .map(_.toMap)
+      override def collectM[G[+_]: IdentityBoth: Covariant, V, V2](map: Map[K, V])(
+        f: V => G[Option[V2]]
+      )(implicit
+        identityBoth: IdentityBoth[({ type lambda[+v] = Map[K, v] })#lambda],
+        identityEither: IdentityEither[({ type lambda[+v] = Map[K, v] })#lambda]
+      ): G[Map[K, V2]] =
+        CovariantIdentityBoth[G]
+          .collectM[(K, V), (K, V2), Iterable](map) { case (k, v) => f(v).map(_.map(k -> _)) }
+          .map(_.toMap)
       override def forEach_[G[+_]: IdentityBoth: Covariant, V](map: Map[K, V])(f: V => G[Any]): G[Unit] =
         CovariantIdentityBoth[G].forEach_(map) { case (_, v) => f(v) }
     }
@@ -1324,6 +1341,10 @@ object Invariant extends LowPriorityInvariantImplicits with InvariantVersionSpec
     new ForEach[Vector] {
       def forEach[G[+_]: IdentityBoth: Covariant, A, B](vector: Vector[A])(f: A => G[B]): G[Vector[B]]     =
         CovariantIdentityBoth[G].forEach(vector)(f)
+      override def collectM[G[+_]: IdentityBoth: Covariant, A, B](vector: Vector[A])(
+        f: A => G[Option[B]]
+      )(implicit identityBoth: IdentityBoth[Vector], identityEither: IdentityEither[Vector]): G[Vector[B]] =
+        CovariantIdentityBoth[G].collectM(vector)(f)
       override def forEach_[G[+_]: IdentityBoth: Covariant, A](vector: Vector[A])(f: A => G[Any]): G[Unit] =
         CovariantIdentityBoth[G].forEach_(vector)(f)
     }

core/shared/src/main/scala/zio/prelude/coherent/coherent.scala

@@ -182,6 +182,11 @@ trait CovariantIdentityBoth[F[+_]] extends Covariant[F] with IdentityBoth[F] { s
   private implicit val covariant: Covariant[F]       = self
   private implicit val identityBoth: IdentityBoth[F] = self
 
+  def collectM[A, B, Collection[+Element] <: Iterable[Element]](in: Collection[A])(f: A => F[Option[B]])(implicit
+    bf: BuildFrom[Collection[A], B, Collection[B]]
+  ): F[Collection[B]] =
+    forEach[A, Option[B], Iterable](in)(f).map(_.flatten).map(bf.fromSpecific(in))
+
   def forEach[A, B, Collection[+Element] <: Iterable[Element]](in: Collection[A])(f: A => F[B])(implicit
     bf: BuildFrom[Collection[A], B, Collection[B]]
   ): F[Collection[B]] =

core/shared/src/main/scala/zio/prelude/fx/ZPure.scala

@@ -991,6 +991,29 @@ object ZPure {
       }
     }
 
+  /**
+   * Evaluate each computation in the structure from left to right, collecting
+   * the successful values and discarding the empty cases.
+   */
+  def collect[W, S, R, E, A, B, Collection[+Element] <: Iterable[Element]](in: Collection[A])(
+    f: A => ZPure[W, S, S, R, E, Option[B]]
+  )(implicit bf: BuildFrom[Collection[A], B, Collection[B]]): ZPure[W, S, S, R, E, Collection[B]] =
+    ZPure.suspend {
+      val iterator = in.iterator
+      val builder  = bf.newBuilder(in)
+
+      lazy val recurse: Option[B] => ZPure[W, S, S, R, E, Collection[B]] = {
+        case Some(b) => builder += b; loop()
+        case None    => loop()
+      }
+
+      def loop(): ZPure[W, S, S, R, E, Collection[B]] =
+        if (iterator.hasNext) f(iterator.next()).flatMap(recurse)
+        else ZPure.succeed(builder.result())
+
+      loop()
+    }
+
   /**
    * Combines a collection of computations into a single computation that
    * passes the updated state from each computation to the next and collects
@@ -1258,6 +1281,10 @@ object ZPure {
         fa.zip(fb)
       def map[A, B](f: A => B): ZPure[W, S, S, R, E, A] => ZPure[W, S, S, R, E, B]                                 =
         _.map(f)
+      override def collectM[A, B, Collection[+Element] <: Iterable[Element]](in: Collection[A])(
+        f: A => ZPure[W, S, S, R, E, Option[B]]
+      )(implicit bf: BuildFrom[Collection[A], B, Collection[B]]): ZPure[W, S, S, R, E, Collection[B]] =
+        ZPure.collect(in)(f)
       override def forEach[A, B, Collection[+Element] <: Iterable[Element]](
         in: Collection[A]
       )(f: A => ZPure[W, S, S, R, E, B])(implicit