Iterate through queried fields instead of object fields (#2021)

core/src/main/scala-2/caliban/schema/SchemaDerivation.scala

@@ -4,7 +4,7 @@ import caliban.Value._
 import caliban.introspection.adt._
 import caliban.parsing.adt.Directive
 import caliban.schema.Annotations._
-import caliban.schema.Step.{ PureStep => _, _ }
+import caliban.schema.Step.{ PureStep => _ }
 import caliban.schema.Types._
 import magnolia1._
 
@@ -36,14 +36,18 @@ trait CommonSchemaDerivation[R] {
     }
 
   def join[T](ctx: ReadOnlyCaseClass[Typeclass, T]): Typeclass[T] = new Typeclass[T] {
-    private lazy val fields = ctx.parameters.map { p =>
-      (getName(p), p.typeclass, p.dereference _)
-    }
+    private lazy val objectResolver =
+      new ObjectFieldResolver[R, T](
+        getName(ctx),
+        ctx.parameters.map { p =>
+          getName(p) -> { (v: T) => p.typeclass.resolve(p.dereference(v)) }
+        }
+      )
 
     private lazy val _isValueType = (ctx.isValueClass || isValueType(ctx)) && ctx.parameters.nonEmpty
 
     override def toType(isInput: Boolean, isSubscription: Boolean): __Type = {
-      val _ = fields // Initializes lazy val
+      val _ = objectResolver // Initializes lazy val
       if (_isValueType) {
         if (isScalarValueType(ctx)) makeScalar(getName(ctx), getDescription(ctx))
         else ctx.parameters.head.typeclass.toType_(isInput, isSubscription)
@@ -98,21 +102,13 @@ trait CommonSchemaDerivation[R] {
     override def resolve(value: T): Step[R] =
       if (ctx.isObject) PureStep(EnumValue(getName(ctx)))
       else if (_isValueType) resolveValueType(value)
-      else MetadataFunctionStep[R](f => resolveObject(value, f.fieldNames))
+      else objectResolver.resolve(value)
 
     private def resolveValueType(value: T): Step[R] = {
       val head = ctx.parameters.head
       head.typeclass.resolve(head.dereference(value))
     }
 
-    private def resolveObject(value: T, queriedFields: Set[String]): Step[R] = {
-      val fieldsBuilder = Map.newBuilder[String, Step[R]]
-      fields.foreach { case (name, schema, dereference) =>
-        if (queriedFields.contains(name))
-          fieldsBuilder += name -> schema.resolve(dereference(value))
-      }
-      ObjectStep(getName(ctx), fieldsBuilder.result())
-    }
   }
 
   def split[T](ctx: SealedTrait[Typeclass, T]): Typeclass[T] = new Typeclass[T] {

core/src/main/scala-3/caliban/schema/ObjectSchema.scala

@@ -2,7 +2,6 @@ package caliban.schema
 
 import caliban.introspection.adt.__Type
 import caliban.schema.DerivationUtils.*
-import caliban.schema.Step.{ MetadataFunctionStep, ObjectStep }
 import magnolia1.TypeInfo
 
 final private class ObjectSchema[R, A](
@@ -11,27 +10,24 @@ final private class ObjectSchema[R, A](
   anns: List[Any],
   paramAnnotations: Map[String, List[Any]]
 ) extends Schema[R, A] {
-  private val name = getName(anns, info)
 
   private lazy val fields = _fields.map { (label, schema, index) =>
     val fieldAnnotations = paramAnnotations.getOrElse(label, Nil)
     (getName(fieldAnnotations, label), fieldAnnotations, schema, index)
   }
 
-  def toType(isInput: Boolean, isSubscription: Boolean): __Type =
-    if (isInput) mkInputObject[R](anns, fields, info)(isInput, isSubscription)
-    else mkObject[R](anns, fields, info)(isInput, isSubscription)
-
-  def resolve(value: A): Step[R] = MetadataFunctionStep[R] { f =>
-    val fb = Map.newBuilder[String, Step[R]]
-
-    var remaining = fields
-    while (!remaining.isEmpty) {
-      val (name, _, schema, i) = remaining.head
-      if (f.fieldNames.contains(name)) fb += name -> schema.resolve(value.asInstanceOf[Product].productElement(i))
-      remaining = remaining.tail
+  private lazy val resolver = {
+    def fs = fields.map { (name, _, schema, i) =>
+      name -> { (v: A) => schema.resolve(v.asInstanceOf[Product].productElement(i)) }
     }
+    new ObjectFieldResolver(getName(anns, info), fs)
+  }
 
-    ObjectStep(name, fb.result())
+  def toType(isInput: Boolean, isSubscription: Boolean): __Type = {
+    val _ = resolver // Init the lazy val
+    if (isInput) mkInputObject[R](anns, fields, info)(isInput, isSubscription)
+    else mkObject[R](anns, fields, info)(isInput, isSubscription)
   }
+
+  def resolve(value: A): Step[R] = resolver.resolve(value)
 }

core/src/main/scala/caliban/execution/Field.scala

@@ -43,11 +43,10 @@ case class Field(
 ) { self =>
   lazy val locationInfo: LocationInfo = _locationInfo()
 
-  private[caliban] lazy val fieldNames: Set[String] =
-    fields.foldLeft(Set.newBuilder[String]) { case (sb, f) => sb += f.name }.result()
-
   private[caliban] val aliasedName: String = alias.getOrElse(name)
 
+  private[caliban] lazy val distinctFieldNames: List[String] = fields.map(_.name).distinct
+
   def combine(other: Field): Field =
     self.copy(
       fields = self.fields ::: other.fields,

core/src/main/scala/caliban/schema/ObjectFieldResolver.scala

@@ -0,0 +1,36 @@
+package caliban.schema
+
+import caliban.execution.Field
+import caliban.schema.Step.{ MetadataFunctionStep, ObjectStep }
+
+import scala.collection.immutable.HashMap
+
+final private class ObjectFieldResolver[R, A](
+  objectName: String,
+  fields: Iterable[(String, A => Step[R])]
+) {
+
+  private val fieldsMap: java.util.HashMap[String, A => Step[R]] = {
+    val map = new java.util.HashMap[String, A => Step[R]]()
+    fields.foreach { case (name, resolve) => map.put(name, resolve) }
+    map
+  }
+
+  def resolve(value: A): Step[R] = MetadataFunctionStep(resolveForField(value, _))
+
+  private def resolveForField(
+    value: A,
+    field: Field
+  ): Step[R] = {
+    val fieldsBuilder = HashMap.newBuilder[String, Step[R]]
+    var remaining     = field.distinctFieldNames
+    while (!remaining.isEmpty) {
+      val name    = remaining.head
+      val resolve = fieldsMap.get(name)
+      if (resolve eq null) ()
+      else fieldsBuilder += name -> resolve(value)
+      remaining = remaining.tail
+    }
+    ObjectStep(objectName, fieldsBuilder.result())
+  }
+}

core/src/main/scala/caliban/schema/Schema.scala

@@ -188,16 +188,10 @@ trait GenericSchema[R] extends SchemaDerivation[R] with TemporalSchema {
           )
         } else makeObject(Some(name), description, fields(isInput, isSubscription).map(_._1), directives)
 
-      private lazy val fieldsForResolve = fields(false, false)
+      private lazy val resolver =
+        new ObjectFieldResolver[R1, A](name, fields(false, false).map(f => (f._1.name, f._2)))
 
-      override def resolve(value: A): Step[R1] = MetadataFunctionStep[R1] { field =>
-        val fieldsBuilder = Map.newBuilder[String, Step[R1]]
-        fieldsForResolve.foreach { case (f, plan) =>
-          if (field.fieldNames.contains(f.name))
-            fieldsBuilder += f.name -> plan(value)
-        }
-        ObjectStep(name, fieldsBuilder.result())
-      }
+      override def resolve(value: A): Step[R1] = resolver.resolve(value)
     }
 
   /**