Add method to memoize the result of a ZQuery

benchmarks/src/main/scala/zio/query/FromRequestBenchmark.scala

@@ -55,6 +55,12 @@ class FromRequestBenchmark {
     unsafeRun(query *> query *> query)
   }
 
+  @Benchmark
+  def fromRequestZipRightMemoized(): Long = {
+    val f = ZQuery.collectAllBatched(queries).map(_.sum.toLong).memoize
+    unsafeRun(ZQuery.unwrap(f.map(q => q *> q *> q)))
+  }
+
   private case class Req(i: Int) extends Request[Nothing, Int]
   private val ds = DataSource.fromFunctionBatchedZIO("Datasource") { reqs: Chunk[Req] => ZIO.succeed(reqs.map(_.i)) }
 }

zio-query/shared/src/main/scala/zio/query/ZQuery.scala

@@ -158,6 +158,9 @@ final class ZQuery[-R, +E, +A] private (private val step: ZIO[R, Nothing, Result
    * Enables caching for this query. Note that caching is enabled by default so
    * this will only be effective to enable caching in part of a larger query in
    * which caching has been disabled.
+   *
+   * @see
+   *   [[memoize]] for memoizing the result of a single query
    */
   def cached(implicit trace: Trace): ZQuery[R, E, A] =
     ZQuery.acquireReleaseWith(ZQuery.cachingEnabled.getAndSet(true))(ZQuery.cachingEnabled.set)(_ => self)
@@ -275,6 +278,35 @@ final class ZQuery[-R, +E, +A] private (private val step: ZIO[R, Nothing, Result
       a => ev(a).fold(b => ZQuery.succeed(b), c => ZQuery.fail(Right(c)))
     )
 
+  /**
+   * Returns an effect that, that if evaluated, will return a lazily computed
+   * version of this query
+   *
+   * This differs from query caching, as caching will only cache the output of a
+   * [[DataSource]]. `memoize` will ensure that the query (including
+   * non-DataSource backed queries) is computed at-most-once.
+   *
+   * This can beneficial for cases that a query is composed of multiple queries
+   * and it's reused multiple times, e.g.,
+   *
+   * {{{
+   *    case class Foo(x: UQuery[Int], y: UQuery[Int])
+   *    val query: UQuery[Int] = ???
+   *
+   *    // Query will be run exactly once; might not be necessary if `x` or `y` are not used afterwards
+   *    query.map(i => Foo(ZQuery.succeed(i +1), ZQuery.succeed(i + 2)))
+   *
+   *    // Query will be recomputed each time x or y are used
+   *    Foo(query.map(_ + 1), query.map(_ + 2))
+   *
+   *    // Query will be computed / run at-most-once
+   *    query.memoize.map(q => Foo(q.map(_ + 1), q.map(_ + 2)))
+   *
+   * }}}
+   */
+  final def memoize(implicit trace: Trace): UIO[ZQuery[R, E, A]] =
+    ZIO.succeed(unsafe.memoize(Unsafe.unsafe, trace))
+
   /**
    * Maps the specified function over the successful result of this query.
    */
@@ -833,6 +865,27 @@ final class ZQuery[-R, +E, +A] private (private val step: ZIO[R, Nothing, Result
         case (_, Result.Fail(e))                                => Result.fail(e)
       }
     }
+
+  /**
+   * These methods can improve UX and performance in some cases, but when used
+   * improperly they can lead to unexpected behaviour in the application code.
+   *
+   * Make sure you really understand them before using them!
+   */
+  def unsafe: UnsafeApi = new UnsafeApi {}
+
+  trait UnsafeApi {
+
+    def memoize(implicit unsafe: Unsafe, trace: Trace): ZQuery[R, E, A] = {
+      val ref = Ref.Synchronized.unsafe.make[Option[Result[R, E, A]]](None)
+      new ZQuery[R, E, A](ref.modifyZIO {
+        case s @ Some(result) => Exit.succeed((result, s))
+        case _                => self.step.map(result => (result, Some(result)))
+      })
+    }
+
+  }
+
 }
 
 object ZQuery {
@@ -1344,7 +1397,7 @@ object ZQuery {
    * one of its variants for optimizations to be applied.
    *
    * @see
-   *   [[fromRequests]] for variants that allow for multiple requests to be
+   *   `fromRequests` for variants that allow for multiple requests to be
    *   submitted at once
    */
   def fromRequest[R, E, A, B](
@@ -1378,8 +1431,8 @@ object ZQuery {
    * @see
    *   [[fromRequest]] for submitting a single request to a datasource
    * @see
-   *   [[fromRequestsWith]] for a variant that allows transforming the input to
-   *   a request
+   *   `fromRequestsWith` for a variant that allows transforming the input to a
+   *   request
    */
   def fromRequests[R, E, A, B](
     requests: Chunk[A]
@@ -1394,8 +1447,8 @@ object ZQuery {
    * @see
    *   [[fromRequest]] for submitting a single request to a datasource
    * @see
-   *   [[fromRequestsWith]] for a variant that allows transforming the input to
-   *   a request
+   *   `fromRequestsWith` for a variant that allows transforming the input to a
+   *   request
    */
   def fromRequests[R, E, A, B](
     requests: List[A]

zio-query/shared/src/test/scala/zio/query/ZQuerySpec.scala

@@ -333,7 +333,44 @@ object ZQuerySpec extends ZIOBaseSpec {
         for {
           _ <- query.run
         } yield assertCompletes
-      }
+      },
+      suite("memoize")(
+        test("results are memoized") {
+          for {
+            ref   <- Ref.make(0)
+            query <- ZQuery.fromZIO(ref.update(_ + 1)).memoize
+            _     <- (ZQuery.foreachPar((1 to 100).toList)(_ => query) <* query).run
+            value <- ref.get
+          } yield assertTrue(value == 1)
+        },
+        test("results are not computed when the outer effect is executed") {
+          for {
+            ref   <- Ref.make(0)
+            query <- ZQuery.fromZIO(ref.update(_ + 1)).memoize
+            value <- ref.get
+          } yield assertTrue(value == 0)
+        },
+        test("errors are memoized") {
+          for {
+            ref     <- Ref.make(0)
+            query   <- ZQuery.fromZIO(ref.updateAndGet(_ + 1).flatMap(i => ZIO.fail(new Exception(i.toString)))).memoize
+            results <- ZQuery.foreachPar((1 to 100).toList)(_ => query.either).run
+            value   <- ref.get
+          } yield assertTrue(value == 1, results.forall(_.isLeft))
+        },
+        test("defects are memoized") {
+          for {
+            ref   <- Ref.make(0)
+            query <- ZQuery.fromZIO(ref.updateAndGet(_ + 1).flatMap(i => ZIO.die(new Exception(i.toString)))).memoize
+            results <- ZQuery
+                         .foreachPar((1 to 100).toList)(_ =>
+                           query.foldCauseQuery(c => ZQuery.succeed(Left(c)), v => ZQuery.succeed(Right(v)))
+                         )
+                         .run
+            value <- ref.get
+          } yield assertTrue(value == 1, results.forall(_.isLeft))
+        }
+      )
     ) @@ silent
 
   val userIds: List[Int]          = (1 to 26).toList