Disable key fields check if unnecessary and optimize its perf

apollo-ast/src/main/kotlin/com/apollographql/apollo3/ast/Schema.kt

@@ -2,7 +2,6 @@ package com.apollographql.apollo3.ast
 
 import com.apollographql.apollo3.annotations.ApolloExperimental
 import com.apollographql.apollo3.ast.internal.buffer
-import okio.Buffer
 
 /**
  * A wrapper around a schema GQLDocument that:
@@ -92,6 +91,15 @@ class Schema(
     }
   }
 
+  /**
+   * Returns whether the `typePolicy` directive is present on at least one object in the schema
+   */
+  fun hasTypeWithTypePolicy(): Boolean {
+    return typeDefinitions.values.filterIsInstance<GQLObjectTypeDefinition>().any { objectType ->
+      objectType.directives.any { it.name == TYPE_POLICY }
+    }
+  }
+
   /**
    * Returns the key fields for the given type
    *

apollo-ast/src/main/kotlin/com/apollographql/apollo3/ast/check_key_fields.kt

@@ -3,14 +3,28 @@ package com.apollographql.apollo3.ast
 private class CheckKeyFieldsScope(
     val schema: Schema,
     val allFragmentDefinitions: Map<String, GQLFragmentDefinition>,
-)
+) {
+  private val implementedTypesCache = mutableMapOf<String, Set<String>>()
+  fun implementedTypes(name: String) = implementedTypesCache.getOrPut(name) {
+    schema.implementedTypes(name)
+  }
+
+  private val keyFieldsCache = mutableMapOf<String, Set<String>>()
+  fun keyFields(name: String) = keyFieldsCache.getOrPut(name) {
+    schema.keyFields(name)
+  }
+}
 
 fun checkKeyFields(operation: GQLOperationDefinition, schema: Schema, allFragmentDefinitions: Map<String, GQLFragmentDefinition>) {
   val parentType = operation.rootTypeDefinition(schema)!!.name
   CheckKeyFieldsScope(schema, allFragmentDefinitions).checkField("Operation(${operation.name})", operation.selectionSet.selections, parentType)
 }
 
-fun checkKeyFields(fragmentDefinition: GQLFragmentDefinition, schema: Schema, allFragmentDefinitions: Map<String, GQLFragmentDefinition>) {
+fun checkKeyFields(
+    fragmentDefinition: GQLFragmentDefinition,
+    schema: Schema,
+    allFragmentDefinitions: Map<String, GQLFragmentDefinition>,
+) {
   CheckKeyFieldsScope(schema, allFragmentDefinitions).checkField("Fragment(${fragmentDefinition.name})", fragmentDefinition.selectionSet.selections, fragmentDefinition.typeCondition.name)
 }
 
@@ -25,9 +39,9 @@ private fun CheckKeyFieldsScope.checkField(
 }
 
 private fun CheckKeyFieldsScope.checkFieldSet(path: String, selections: List<GQLSelection>, parentType: String, possibleType: String) {
-  val implementedTypes = schema.implementedTypes(possibleType)
+  val implementedTypes = implementedTypes(possibleType)
 
-  val mergedFields = collectFields(selections, parentType, implementedTypes).groupBy {
+  val mergedFields = collectFieldsCached(selections, parentType, implementedTypes).groupBy {
     it.field.name
   }.values
 
@@ -36,7 +50,7 @@ private fun CheckKeyFieldsScope.checkFieldSet(path: String, selections: List<GQL
     val fieldNames = mergedFields.map { it.first().field }
         .filter { it.alias == null }
         .map { it.name }.toSet()
-    val keyFieldNames = schema.keyFields(possibleType)
+    val keyFieldNames = keyFields(possibleType)
 
     val missingFieldNames = keyFieldNames.subtract(fieldNames)
     check(missingFieldNames.isEmpty()) {
@@ -53,6 +67,15 @@ private fun CheckKeyFieldsScope.checkFieldSet(path: String, selections: List<GQL
 
 private class FieldWithParent(val field: GQLField, val parentType: String)
 
+private val collectFieldsCache = mutableMapOf<String, List<FieldWithParent>>()
+private fun CheckKeyFieldsScope.collectFieldsCached(
+    selections: List<GQLSelection>,
+    parentType: String,
+    implementedTypes: Set<String>,
+) = collectFieldsCache.getOrPut("$selections $parentType $implementedTypes") {
+  collectFields(selections, parentType, implementedTypes)
+}
+
 private fun CheckKeyFieldsScope.collectFields(
     selections: List<GQLSelection>,
     parentType: String,
@@ -75,15 +98,15 @@ private fun CheckKeyFieldsScope.collectFields(
           return@flatMap emptyList()
         }
 
-        collectFields(it.selectionSet.selections, it.typeCondition.name, implementedTypes)
+        collectFieldsCached(it.selectionSet.selections, it.typeCondition.name, implementedTypes)
       }
       is GQLFragmentSpread -> {
         if (it.directives.hasCondition()) {
           return@flatMap emptyList()
         }
 
         val fragmentDefinition = allFragmentDefinitions[it.name]!!
-        collectFields(fragmentDefinition.selectionSet.selections, fragmentDefinition.typeCondition.name, implementedTypes)
+        collectFieldsCached(fragmentDefinition.selectionSet.selections, fragmentDefinition.typeCondition.name, implementedTypes)
       }
     }
   }
@@ -94,4 +117,4 @@ private fun List<GQLDirective>?.hasCondition(): Boolean {
     it.name == "skip" && (it.arguments!!.arguments.first().value as? GQLStringValue)?.value != "false"
         || it.name == "include" && (it.arguments!!.arguments.first().value as? GQLStringValue)?.value != "true"
   } ?: false
-}
+}

apollo-compiler/src/main/kotlin/com/apollographql/apollo3/compiler/ApolloCompiler.kt

@@ -116,11 +116,15 @@ object ApolloCompiler {
     // Remember the fragments with the possibly updated fragments
     val allFragmentDefinitions = (fragments + incomingFragments).associateBy { it.name }
 
-    operations.forEach {
-      checkKeyFields(it, options.schema, allFragmentDefinitions)
-    }
-    fragments.forEach {
-      checkKeyFields(it, options.schema, allFragmentDefinitions)
+    // Check if all the key fields are present in operations and fragments
+    // (do this only if there are key fields as it may be costly)
+    if (options.schema.hasTypeWithTypePolicy()) {
+      operations.forEach {
+        checkKeyFields(it, options.schema, allFragmentDefinitions)
+      }
+      fragments.forEach {
+        checkKeyFields(it, options.schema, allFragmentDefinitions)
+      }
     }
 
     var alwaysGenerateTypesMatching = options.alwaysGenerateTypesMatching