Implement ObjectStep using a Function1 (#2031)

* Implement ObjectStep using a Function1

* Add collection compat import for Scala 2.12

* Resolve unknown fields to `NullStep`

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

@@ -37,7 +37,7 @@ trait CommonSchemaDerivation[R] {
 
   def join[T](ctx: ReadOnlyCaseClass[Typeclass, T]): Typeclass[T] = new Typeclass[T] {
     private lazy val objectResolver =
-      new ObjectFieldResolver[R, T](
+      ObjectFieldResolver[R, T](
         getName(ctx),
         ctx.parameters.map { p =>
           getName(p) -> { (v: T) => p.typeclass.resolve(p.dereference(v)) }

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

@@ -20,7 +20,7 @@ final private class ObjectSchema[R, A](
     def fs = fields.map { (name, _, schema, i) =>
       name -> { (v: A) => schema.resolve(v.asInstanceOf[Product].productElement(i)) }
     }
-    new ObjectFieldResolver(getName(anns, info), fs)
+    ObjectFieldResolver(getName(anns, info), fs)
   }
 
   def toType(isInput: Boolean, isSubscription: Boolean): __Type = {

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

@@ -59,16 +59,13 @@ object Executor {
       path: List[Either[String, Int]]
     ): ReducedStep[R] = {
 
-      def reduceObjectStep(objectName: String, fields: collection.Map[String, Step[R]]): ReducedStep[R] = {
+      def reduceObjectStep(objectName: String, getFieldStep: String => Step[R]): ReducedStep[R] = {
         val filteredFields    = mergeFields(currentField, objectName)
         val (deferred, eager) = filteredFields.partitionMap {
           case f @ Field("__typename", _, _, _, _, _, _, directives, _, _, _)   =>
             Right((f.aliasedName, PureStep(StringValue(objectName)), fieldInfo(f, path, directives)))
           case f @ Field(name, _, _, _, _, _, args, directives, _, _, fragment) =>
-            val field = fields.get(name) match {
-              case Some(step) => reduceStep(step, f, args, Left(f.aliasedName) :: path)
-              case _          => NullStep
-            }
+            val field = reduceStep(getFieldStep(name), f, args, Left(f.aliasedName) :: path)
             val entry = (f.aliasedName, field, fieldInfo(f, path, directives))
 
             fragment match {

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

@@ -45,8 +45,6 @@ case class Field(
 
   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

@@ -1,36 +1,25 @@
 package caliban.schema
 
-import caliban.execution.Field
-import caliban.schema.Step.{ MetadataFunctionStep, ObjectStep }
+import caliban.schema.Step.{ NullStep, ObjectStep }
 
+import scala.collection.compat._
 import scala.collection.mutable
 
-final private class ObjectFieldResolver[R, A](
-  objectName: String,
-  fields: Iterable[(String, A => Step[R])]
+final private class ObjectFieldResolver[R, A] private (
+  name: String,
+  fields: mutable.HashMap[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 = new mutable.HashMap[String, Step[R]]()
-
-    var remaining = field.distinctFieldNames
-    while (!remaining.isEmpty) {
-      val name    = remaining.head
-      val resolve = fieldsMap.get(name)
-      if (resolve ne null) fieldsBuilder.update(name, resolve(value))
-      remaining = remaining.tail
+  def resolve(value: A): Step[R] = ObjectStep(
+    name,
+    fields.get(_) match {
+      case Some(f) => f(value)
+      case None    => NullStep
     }
-    ObjectStep(objectName, fieldsBuilder)
-  }
+  )
+}
+
+private object ObjectFieldResolver {
+  def apply[R, A](objectName: String, fields: Iterable[(String, A => Step[R])]): ObjectFieldResolver[R, A] =
+    // NOTE: mutable.HashMap is about twice as fast than immutable.HashMap for .get
+    new ObjectFieldResolver(objectName, mutable.HashMap.from(fields))
 }

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

@@ -176,7 +176,8 @@ trait GenericSchema[R] extends SchemaDerivation[R] with TemporalSchema {
     directives: List[Directive] = List.empty
   ): Schema[R1, A] =
     new Schema[R1, A] {
-      override def toType(isInput: Boolean, isSubscription: Boolean): __Type =
+      override def toType(isInput: Boolean, isSubscription: Boolean): __Type = {
+        val _ = resolver
         if (isInput) {
           makeInputObject(
             Some(customizeInputTypeName(name)),
@@ -187,9 +188,10 @@ trait GenericSchema[R] extends SchemaDerivation[R] with TemporalSchema {
             directives = Some(directives)
           )
         } else makeObject(Some(name), description, fields(isInput, isSubscription).map(_._1), directives)
+      }
 
       private lazy val resolver =
-        new ObjectFieldResolver[R1, A](name, fields(false, false).map(f => (f._1.name, f._2)))
+        ObjectFieldResolver[R1, A](name, fields(false, false).map(f => (f._1.name, f._2)))
 
       override def resolve(value: A): Step[R1] = resolver.resolve(value)
     }

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

@@ -10,12 +10,17 @@ import zio.stream.ZStream
 sealed trait Step[-R]
 
 object Step {
-  case class ListStep[-R](steps: List[Step[R]])                                    extends Step[R]
-  case class FunctionStep[-R](step: Map[String, InputValue] => Step[R])            extends Step[R]
-  case class MetadataFunctionStep[-R](step: Field => Step[R])                      extends Step[R]
-  case class ObjectStep[-R](name: String, fields: collection.Map[String, Step[R]]) extends Step[R]
-  case class QueryStep[-R](query: ZQuery[R, Throwable, Step[R]])                   extends Step[R]
-  case class StreamStep[-R](inner: ZStream[R, Throwable, Step[R]])                 extends Step[R]
+  case class ListStep[-R](steps: List[Step[R]])                         extends Step[R]
+  case class FunctionStep[-R](step: Map[String, InputValue] => Step[R]) extends Step[R]
+  case class MetadataFunctionStep[-R](step: Field => Step[R])           extends Step[R]
+  case class QueryStep[-R](query: ZQuery[R, Throwable, Step[R]])        extends Step[R]
+  case class StreamStep[-R](inner: ZStream[R, Throwable, Step[R]])      extends Step[R]
+
+  case class ObjectStep[-R](name: String, fields: String => Step[R]) extends Step[R]
+  object ObjectStep {
+    def apply[R](name: String, fields: Map[String, Step[R]]): ObjectStep[R] =
+      new ObjectStep[R](name, fields.getOrElse(_, NullStep))
+  }
 
   // PureStep is both a Step and a ReducedStep so it is defined outside this object
   // This is to avoid boxing/unboxing pure values during step reduction
@@ -34,13 +39,20 @@ object Step {
     case (FunctionStep(l), FunctionStep(r))                  => FunctionStep(args => mergeRootSteps(l(args), r(args)))
     case (FunctionStep(l), r)                                => FunctionStep(args => mergeRootSteps(l(args), r))
     case (l, FunctionStep(r))                                => FunctionStep(args => mergeRootSteps(l, r(args)))
-    // fields2 override fields1 in case of conflict
-    case (ObjectStep(name, fields1), ObjectStep(_, fields2)) => ObjectStep(name, fields1 ++ fields2)
+    case (ObjectStep(name, fields1), ObjectStep(_, fields2)) => ObjectStep(name, mergeObjectSteps(fields1, fields2))
     // if only step1 is an object, keep it
     case (ObjectStep(_, _), _)                               => step1
     // otherwise keep step2
     case _                                                   => step2
   }
+
+  // fields2 override fields1 in case of conflict
+  private def mergeObjectSteps[R](fields1: String => Step[R], fields2: String => Step[R]): String => Step[R] =
+    (s: String) =>
+      fields2(s) match {
+        case NullStep => fields1(s)
+        case step     => step
+      }
 }
 
 sealed trait ReducedStep[-R] { self =>

interop/cats/src/test/scala/caliban/interop/fs2/Fs2InteropSchemaSpec.scala

@@ -7,12 +7,11 @@ import caliban.introspection.adt.{ __Type, __TypeKind }
 import caliban.parsing.adt.Definition.TypeSystemDefinition.TypeDefinition
 import caliban.parsing.adt.Definition.TypeSystemDefinition.TypeDefinition.{ FieldDefinition, ObjectTypeDefinition }
 import caliban.parsing.adt.Type._
-import caliban.schema.Schema.auto._
 import caliban.schema.Step.{ MetadataFunctionStep, ObjectStep, StreamStep }
 import caliban.schema.{ PureStep, Schema, Step }
 import cats.effect.std.Dispatcher
 import cats.effect.unsafe.implicits._
-import cats.effect.{ IO, LiftIO, Resource }
+import cats.effect.{ IO, LiftIO }
 import fs2.Stream
 import zio.interop.catz._
 import zio.test.Assertion._
@@ -120,19 +119,16 @@ object Fs2InteropSchemaSpec extends ZIOSpecDefault {
     isSubtype[ObjectStep[Any]](
       hasField(
         "fields",
-        _.fields.toMap,
-        hasKey(
-          "bar",
-          isSubtype[StreamStep[Any]](
-            hasField(
-              "inner",
-              step =>
-                Unsafe.unsafe { implicit unsafe =>
-                  runtime.unsafe.run(step.inner.runCollect)
-                },
-              isSubtype[Exit.Success[Chunk[Step[Any]]]](
-                equalTo(Exit.Success(expectedChunk))
-              )
+        _.fields("bar"),
+        isSubtype[StreamStep[Any]](
+          hasField(
+            "inner",
+            step =>
+              Unsafe.unsafe { implicit unsafe =>
+                runtime.unsafe.run(step.inner.runCollect)
+              },
+            isSubtype[Exit.Success[Chunk[Step[Any]]]](
+              equalTo(Exit.Success(expectedChunk))
             )
           )
         )