From d5bf339cf73003a35c474483046e9ea3ce06ad6e Mon Sep 17 00:00:00 2001 From: Adam Fraser Date: Tue, 9 Jan 2024 15:50:53 -0800 Subject: [PATCH 1/2] update signature of foreach collectm --- .../scala/zio/prelude/CollectBenchmark.scala | 53 ++++++++++++++++ .../test/scala/zio/prelude/ForEachSpec.scala | 11 ++++ .../prelude/InvariantVersionSpecific.scala | 4 ++ .../prelude/InvariantVersionSpecific.scala | 4 ++ .../scala/zio/prelude/AssociativeBoth.scala | 14 +++++ .../src/main/scala/zio/prelude/ForEach.scala | 62 +++++++------------ .../main/scala/zio/prelude/Invariant.scala | 21 +++++++ .../scala/zio/prelude/coherent/coherent.scala | 5 ++ .../src/main/scala/zio/prelude/fx/ZPure.scala | 27 ++++++++ 9 files changed, 160 insertions(+), 41 deletions(-) create mode 100644 benchmarks/src/main/scala/zio/prelude/CollectBenchmark.scala diff --git a/benchmarks/src/main/scala/zio/prelude/CollectBenchmark.scala b/benchmarks/src/main/scala/zio/prelude/CollectBenchmark.scala new file mode 100644 index 000000000..efca6b590 --- /dev/null +++ b/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()) +} diff --git a/core-tests/shared/src/test/scala/zio/prelude/ForEachSpec.scala b/core-tests/shared/src/test/scala/zio/prelude/ForEachSpec.scala index 606f689ed..b80d72d3a 100644 --- a/core-tests/shared/src/test/scala/zio/prelude/ForEachSpec.scala +++ b/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,14 @@ 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) diff --git a/core/shared/src/main/scala-2.12/zio/prelude/InvariantVersionSpecific.scala b/core/shared/src/main/scala-2.12/zio/prelude/InvariantVersionSpecific.scala index 6c19e5e73..aacbf9245 100644 --- a/core/shared/src/main/scala-2.12/zio/prelude/InvariantVersionSpecific.scala +++ b/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) } diff --git a/core/shared/src/main/scala-2.13+/zio/prelude/InvariantVersionSpecific.scala b/core/shared/src/main/scala-2.13+/zio/prelude/InvariantVersionSpecific.scala index f14acef29..3c6d17ff9 100644 --- a/core/shared/src/main/scala-2.13+/zio/prelude/InvariantVersionSpecific.scala +++ b/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) } diff --git a/core/shared/src/main/scala/zio/prelude/AssociativeBoth.scala b/core/shared/src/main/scala/zio/prelude/AssociativeBoth.scala index fcd9d82e1..e4a03793d 100644 --- a/core/shared/src/main/scala/zio/prelude/AssociativeBoth.scala +++ b/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]] = diff --git a/core/shared/src/main/scala/zio/prelude/ForEach.scala b/core/shared/src/main/scala/zio/prelude/ForEach.scala index 479870861..9fe3028ff 100644 --- a/core/shared/src/main/scala/zio/prelude/ForEach.scala +++ 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,7 +316,7 @@ 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 @@ -344,8 +324,8 @@ trait ForEach[F[+_]] extends Covariant[F] { self => 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) diff --git a/core/shared/src/main/scala/zio/prelude/Invariant.scala b/core/shared/src/main/scala/zio/prelude/Invariant.scala index e36e297cc..530087664 100644 --- a/core/shared/src/main/scala/zio/prelude/Invariant.scala +++ b/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) } diff --git a/core/shared/src/main/scala/zio/prelude/coherent/coherent.scala b/core/shared/src/main/scala/zio/prelude/coherent/coherent.scala index 7ab57335b..556a3c933 100644 --- a/core/shared/src/main/scala/zio/prelude/coherent/coherent.scala +++ b/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]] = diff --git a/core/shared/src/main/scala/zio/prelude/fx/ZPure.scala b/core/shared/src/main/scala/zio/prelude/fx/ZPure.scala index c2cd1a3a2..d9ed787f0 100644 --- a/core/shared/src/main/scala/zio/prelude/fx/ZPure.scala +++ 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 From 92c0f19376bf34ed022b132494fb4c967150dca0 Mon Sep 17 00:00:00 2001 From: Adam Fraser Date: Tue, 9 Jan 2024 15:58:26 -0800 Subject: [PATCH 2/2] format --- core-tests/shared/src/test/scala/zio/prelude/ForEachSpec.scala | 1 - 1 file changed, 1 deletion(-) diff --git a/core-tests/shared/src/test/scala/zio/prelude/ForEachSpec.scala b/core-tests/shared/src/test/scala/zio/prelude/ForEachSpec.scala index b80d72d3a..8d60d590d 100644 --- a/core-tests/shared/src/test/scala/zio/prelude/ForEachSpec.scala +++ b/core-tests/shared/src/test/scala/zio/prelude/ForEachSpec.scala @@ -59,7 +59,6 @@ object ForEachSpec extends ZIOBaseSpec { 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))