From 8ef7fb3ce54e3b380743b5b61fd5f414177ab4b1 Mon Sep 17 00:00:00 2001
From: kyri-petrou <67301607+kyri-petrou@users.noreply.github.com>
Date: Wed, 22 Nov 2023 18:11:11 +1100
Subject: [PATCH] Don't include pure object fields in ZQuery collection (#2015)

* Don't include pure values in ZQuery collection

* Fix for Scala 2.12

* Go back to using mutable.HashMap
---
 .../NestedZQueryBenchmarkSchema.scala         |  8 ++--
 .../scala/caliban/execution/Executor.scala    | 43 +++++++++++++++----
 2 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/benchmarks/src/main/scala/caliban/execution/NestedZQueryBenchmarkSchema.scala b/benchmarks/src/main/scala/caliban/execution/NestedZQueryBenchmarkSchema.scala
index f9501602cc..c3df6f5861 100644
--- a/benchmarks/src/main/scala/caliban/execution/NestedZQueryBenchmarkSchema.scala
+++ b/benchmarks/src/main/scala/caliban/execution/NestedZQueryBenchmarkSchema.scala
@@ -50,8 +50,8 @@ object NestedZQueryBenchmarkSchema {
   case class MultifieldRoot(entities: Query[List[MultifieldEntity]])
   case class MultifieldEntity(
     id: Int,
-    nested0: Query[Int],
-    nested1: Query[Int],
+    nested0: Int,
+    nested1: Int,
     nested2: Query[Int],
     nested3: Query[Int],
     nested4: Query[Int],
@@ -167,8 +167,8 @@ object NestedZQueryBenchmarkSchema {
       val qi = ZQuery.succeed(i)
       MultifieldEntity(
         i,
-        qi,
-        qi,
+        i + 1,
+        i + 2,
         qi,
         qi,
         qi,
diff --git a/core/src/main/scala/caliban/execution/Executor.scala b/core/src/main/scala/caliban/execution/Executor.scala
index e71880e676..1896e9fdc5 100644
--- a/core/src/main/scala/caliban/execution/Executor.scala
+++ b/core/src/main/scala/caliban/execution/Executor.scala
@@ -15,6 +15,7 @@ import zio.query.{ Cache, ZQuery }
 import zio.stream.ZStream
 
 import scala.annotation.tailrec
+import scala.collection.mutable
 import scala.jdk.CollectionConverters._
 
 object Executor {
@@ -157,20 +158,46 @@ object Executor {
               val q = if (isPure && !wrapper.wrapPureValues) query else wrapper.wrap(query, fieldInfo)
               loop(q, tail)
           }
-        if (isPure && !wrapPureValues) query
-        else loop(query, wrappers)
+        loop(query, wrappers)
+      }
+
+      def objectFieldQuery(name: String, step: ReducedStep[R], info: FieldInfo) = {
+        val q = wrap(loop(step), step.isPure)(fieldWrappers, info)
+        if (info.details.fieldType.isNullable) q.catchAll(handleError).map((name, _))
+        else q.map((name, _))
       }
 
       def loop(step: ReducedStep[R]): ZQuery[R, ExecutionError, ResponseValue] =
         step match {
           case PureStep(value)                               => ZQuery.succeed(value)
           case ReducedStep.ObjectStep(steps)                 =>
-            val queries = steps.map { case (name, step, info) =>
-              val q = wrap(loop(step), step.isPure)(fieldWrappers, info)
-              (if (info.details.fieldType.isNullable) q.catchAll(handleError) else q)
-                .map(name -> _)
-            }
-            collectAll(queries).map(ObjectValue.apply)
+            var resolved: mutable.HashMap[String, ResponseValue] = null
+
+            val queries =
+              if (wrapPureValues) steps.map((objectFieldQuery _).tupled)
+              else {
+                val queries = List.newBuilder[ZQuery[R, ExecutionError, (String, ResponseValue)]]
+
+                var remaining = steps
+                while (!remaining.isEmpty) {
+                  remaining.head match {
+                    case (name, PureStep(value), _) =>
+                      if (null == resolved) resolved = new mutable.HashMap[String, ResponseValue]()
+                      resolved.update(name, value)
+                    case (name, step, info)         =>
+                      queries += objectFieldQuery(name, step, info)
+                  }
+                  remaining = remaining.tail
+                }
+                queries.result()
+              }
+
+            if (null == resolved) collectAll(queries).map(ObjectValue.apply)
+            else
+              collectAll(queries).map { results =>
+                results.foreach(kv => resolved.update(kv._1, kv._2))
+                ObjectValue(steps.map { case (name, _, _) => name -> resolved(name) })
+              }
           case ReducedStep.QueryStep(step)                   =>
             step.flatMap(loop)
           case ReducedStep.ListStep(steps, areItemsNullable) =>