feat: support costs on arguments of directives (#1465)

I have opted for the simple base layer without recursion. 
Directives could have input objects passed to arguments, 
where fields could have costs. 

We do not handle such recursion here. 
I settled on the arguments level of directives.

Author: ysmolski

execution/engine/execution_engine_cost_test.go

@@ -4671,7 +4671,7 @@ func TestExecutionEngine_Execute(t *testing.T) {
 			var ds1CostConfig *plan.DataSourceCostConfig
 			if opts.includeCostConfig {
 				ds1CostConfig = &plan.DataSourceCostConfig{
-					Weights: map[plan.FieldCoordinate]*plan.FieldWeight{
+					Weights: map[plan.FieldCoordinate]*plan.FieldCost{
 						{TypeName: "Query", FieldName: "accounts"}: {HasWeight: true, Weight: 5},
 						{TypeName: "User", FieldName: "some"}:      {HasWeight: true, Weight: 2},
 						{TypeName: "Admin", FieldName: "some"}:     {HasWeight: true, Weight: 3},
@@ -4686,7 +4686,7 @@ func TestExecutionEngine_Execute(t *testing.T) {
 			var ds2CostConfig *plan.DataSourceCostConfig
 			if opts.includeCostConfig {
 				ds2CostConfig = &plan.DataSourceCostConfig{
-					Weights: map[plan.FieldCoordinate]*plan.FieldWeight{
+					Weights: map[plan.FieldCoordinate]*plan.FieldCost{
 						{TypeName: "User", FieldName: "name"}:       {HasWeight: true, Weight: 2},
 						{TypeName: "User", FieldName: "title"}:      {HasWeight: true, Weight: 4},
 						{TypeName: "Admin", FieldName: "adminName"}: {HasWeight: true, Weight: 3},

execution/engine/execution_engine_test.go

@@ -641,6 +641,25 @@
             }
           ]
         },
+        {
+          "name": "approx",
+          "description": "",
+          "locations": [
+            "FIELD_DEFINITION"
+          ],
+          "args": [
+            {
+              "name": "tolerance",
+              "description": "",
+              "type": {
+                "kind": "SCALAR",
+                "name": "Int",
+                "ofType": null
+              },
+              "defaultValue": "1"
+            }
+          ]
+        },
         {
           "name": "include",
           "description": "Directs the executor to include this field or fragment only when the argument is true.",

execution/engine/testdata/full_introspection.json

@@ -675,6 +675,25 @@
             }
           ]
         },
+        {
+          "name": "approx",
+          "description": "",
+          "locations": [
+            "FIELD_DEFINITION"
+          ],
+          "args": [
+            {
+              "name": "tolerance",
+              "description": "",
+              "type": {
+                "kind": "SCALAR",
+                "name": "Int",
+                "ofType": null
+              },
+              "defaultValue": "1"
+            }
+          ]
+        },
         {
           "name": "include",
           "description": "Directs the executor to include this field or fragment only when the argument is true.",

execution/engine/testdata/full_introspection_with_deprecated.json

@@ -725,6 +725,28 @@
             }
           ]
         },
+        {
+          "__typename": "__Directive",
+          "name": "approx",
+          "description": "",
+          "locations": [
+            "FIELD_DEFINITION"
+          ],
+          "args": [
+            {
+              "__typename": "__InputValue",
+              "name": "tolerance",
+              "description": "",
+              "type": {
+                "__typename": "__Type",
+                "kind": "SCALAR",
+                "name": "Int",
+                "ofType": null
+              },
+              "defaultValue": "1"
+            }
+          ]
+        },
         {
           "__typename": "__Directive",
           "name": "include",

execution/engine/testdata/full_introspection_with_typenames.json

@@ -14,11 +14,12 @@ https://ibm.github.io/graphql-specs/cost-spec.html
 
 It builds on top of IBM spec for @cost and @listSize directive with a few changes.
 
-* We use Int! for weights instead of floats packed in String!.
+* We use the Int! type for weights.
 * When weight is specified for the type and a field returns the list of that type,
 this weight (along with children's costs) is multiplied too.
 
-TODO: Weights on arguments of directives
+Weights on arguments of directives are supported. If an argument is of InputObject's type,
+then the weight from its fields is not counted.
 
 */
 
@@ -40,8 +41,8 @@ import (
 const DefaultEnumScalarWeight = 0
 const DefaultObjectWeight = 1
 
-// FieldWeight defines cost configuration for a specific field of an object or input object.
-type FieldWeight struct {
+// FieldCost defines cost configuration for a specific field of an object or input object.
+type FieldCost struct {
 
 	// Weight is the cost of this field definition. It could be negative or zero.
 	// Should be used only if HasWeight is true.
@@ -53,6 +54,10 @@ type FieldWeight struct {
 	// ArgumentWeights maps an argument name to its weight.
 	// Location: ARGUMENT_DEFINITION
 	ArgumentWeights map[string]int
+
+	// DirectiveArgumentWeights maps a directive.argument coords to its weight.
+	// Populated by composition from @cost on directive argument definitions.
+	DirectiveArgumentWeights map[string]int
 }
 
 // FieldListSize contains parsed data from the @listSize directive for an object field.
@@ -118,7 +123,7 @@ func (ls *FieldListSize) multiplier(arguments map[string]ArgumentInfo, vars *ast
 type DataSourceCostConfig struct {
 	// Weights maps field coordinate to its weights. Cannot be on fields of interfaces.
 	// Location: FIELD_DEFINITION, INPUT_FIELD_DEFINITION
-	Weights map[FieldCoordinate]*FieldWeight
+	Weights map[FieldCoordinate]*FieldCost
 
 	// ListSizes maps field coordinates to their respective list size configurations.
 	// Location: FIELD_DEFINITION
@@ -129,19 +134,12 @@ type DataSourceCostConfig struct {
 	// Weight assigned to the field or argument definitions overrides the weight of type definition.
 	// Location: ENUM, OBJECT, SCALAR
 	Types map[string]int
-
-	// Arguments on directives is a special case. They use a special kind of coordinate:
-	// directive name + argument name. That should be the key mapped to the weight.
-	//
-	// Directives can be used on [input] object fields and arguments of fields. This creates
-	// mutual recursion between them; it complicates cost calculation.
-	// We avoid them intentionally in the first iteration.
 }
 
 // NewDataSourceCostConfig creates a new cost config with defaults
 func NewDataSourceCostConfig() *DataSourceCostConfig {
 	return &DataSourceCostConfig{
-		Weights:   make(map[FieldCoordinate]*FieldWeight),
+		Weights:   make(map[FieldCoordinate]*FieldCost),
 		ListSizes: make(map[FieldCoordinate]*FieldListSize),
 		Types:     make(map[string]int),
 	}
@@ -238,7 +236,7 @@ type inputObjectField struct {
 
 // inputFieldsCost computes the cost of input object fields from the variable value.
 // It handles both single objects and arrays of objects.
-func (arg *ArgumentInfo) inputFieldsCost(variables *astjson.Value, weights map[FieldCoordinate]*FieldWeight) int {
+func (arg *ArgumentInfo) inputFieldsCost(variables *astjson.Value, weights map[FieldCoordinate]*FieldCost) int {
 	if !arg.hasVariable {
 		return 0
 	}
@@ -262,8 +260,8 @@ func (arg *ArgumentInfo) inputFieldsCost(variables *astjson.Value, weights map[F
 	return 0
 }
 
-func (node *CostTreeNode) maxWeightImplementingField(config *DataSourceCostConfig, fieldName string) *FieldWeight {
-	var maxWeight *FieldWeight
+func (node *CostTreeNode) maxWeightImplementingField(config *DataSourceCostConfig, fieldName string) *FieldCost {
+	var maxWeight *FieldCost
 	for _, implTypeName := range node.implementingTypeNames {
 		// Get the cost config for the field of an implementing type.
 		coord := FieldCoordinate{implTypeName, fieldName}
@@ -330,6 +328,29 @@ func (node *CostTreeNode) sizedFieldImplementingFields(config *DataSourceCostCon
 	return result
 }
 
+// maxDirectiveArgumentWeightsImplementingFields returns the union of DirectiveArgumentWeights
+// from implementing types' field definitions. For each directive.argument pair, it takes the
+// maximum weight across all implementing types.
+func (node *CostTreeNode) maxDirectiveArgumentWeightsImplementingFields(config *DataSourceCostConfig, fieldName string) map[string]int {
+	var result map[string]int
+	for _, implTypeName := range node.implementingTypeNames {
+		coords := FieldCoordinate{implTypeName, fieldName}
+		fw := config.Weights[coords]
+		if fw == nil || len(fw.DirectiveArgumentWeights) == 0 {
+			continue
+		}
+		if result == nil {
+			result = make(map[string]int)
+		}
+		for dirArg, weight := range fw.DirectiveArgumentWeights {
+			if existing, ok := result[dirArg]; !ok || weight > existing {
+				result[dirArg] = weight
+			}
+		}
+	}
+	return result
+}
+
 // cost calculates the estimated/actual cost of this node and all descendants.
 //
 // defaultListSize designates the mode of operation.
@@ -382,7 +403,7 @@ func (node *CostTreeNode) cost(configs map[DSHash]*DataSourceCostConfig, variabl
 //
 // fieldCost is the weight of this field or its returned type
 // argsCost is the sum of argument weights and input fields used on this field.
-// Weights on directives ignored for now.
+// directiveCost is the sum of directive argument weights.
 //
 // defaultListSize designates the mode of operation.
 // When it is positive, then its value is used as a fallback value of list sizes for the estimated cost.
@@ -391,7 +412,7 @@ func (node *CostTreeNode) cost(configs map[DSHash]*DataSourceCostConfig, variabl
 // When estimating cost, it picks the highest multiplier among different data sources.
 // Also, it picks the maximum field weight of implementing types and then
 // the maximum among slicing arguments.
-func (node *CostTreeNode) costsAndMultiplier(configs map[DSHash]*DataSourceCostConfig, variables *astjson.Value, defaultListSize int, actualListSizes map[string]int) (fieldCost, argsCost, directiveCost int, multiplier float64) {
+func (node *CostTreeNode) costsAndMultiplier(configs map[DSHash]*DataSourceCostConfig, variables *astjson.Value, defaultListSize int, actualListSizes map[string]int) (fieldCost, argsCost, directivesCost int, multiplier float64) {
 	if len(node.dataSourceHashes) <= 0 {
 		// no data source is responsible for this field
 		return
@@ -400,7 +421,7 @@ func (node *CostTreeNode) costsAndMultiplier(configs map[DSHash]*DataSourceCostC
 	parent := node.parent
 	fieldCost = 0
 	argsCost = 0
-	directiveCost = 0
+	directivesCost = 0
 	multiplier = 0
 
 	isEstimation := defaultListSize > 0
@@ -479,6 +500,18 @@ func (node *CostTreeNode) costsAndMultiplier(configs map[DSHash]*DataSourceCostC
 			}
 		}
 
+		// Directive weights: sum from the field's own DirectiveArgumentWeights,
+		// or from implementing types when the enclosing type is abstract.
+		if node.isEnclosingTypeAbstract && parent.returnsAbstractType {
+			for _, weight := range parent.maxDirectiveArgumentWeightsImplementingFields(dsCostConfig, node.fieldCoords.FieldName) {
+				directivesCost += weight
+			}
+		} else if fieldWeight != nil {
+			for _, weight := range fieldWeight.DirectiveArgumentWeights {
+				directivesCost += weight
+			}
+		}
+
 		if !node.returnsListType || !isEstimation {
 			continue
 		}
@@ -571,7 +604,7 @@ func (node *CostTreeNode) costsAndMultiplier(configs map[DSHash]*DataSourceCostC
 func inputObjectCost(
 	typeName string,
 	value *astjson.Object,
-	weights map[FieldCoordinate]*FieldWeight,
+	weights map[FieldCoordinate]*FieldCost,
 	types map[FieldCoordinate]inputObjectField) int {
 	if value == nil {
 		return 0

v2/pkg/engine/plan/cost.go

@@ -101,7 +101,7 @@ func (v *CostVisitor) EnterField(fieldRef int) {
 	}
 
 	isEnclosingTypeAbstract := v.Walker.EnclosingTypeDefinition.Kind.IsAbstractType()
-	// Create a skeleton node. dataSourceHashes will be filled in leaveFieldCost
+	// Partially filled node. dataSourceHashes will be filled in leaveFieldCost
 	node := CostTreeNode{
 		fieldRef:                fieldRef,
 		fieldCoords:             FieldCoordinate{typeName, fieldName},
@@ -115,7 +115,7 @@ func (v *CostVisitor) EnterField(fieldRef int) {
 		jsonPath:                jsonPath,
 	}
 
-	// Attach to parent
+	// Attach to the parent
 	if len(v.stack) > 0 {
 		parent := v.stack[len(v.stack)-1]
 		parent.children = append(parent.children, &node)

v2/pkg/engine/plan/cost_visitor.go

@@ -8,10 +8,14 @@ schema {
 
 directive @testDeprecated(okArg: String deprecatedArg: String @deprecated(reason: "no such arg")) on FIELD_DEFINITION
 
+# Used to test costs on arguments of directives:
+# the tolerance argument has "@cost(weight: -5)" defined in tests.
+directive @approx(tolerance: Int = 1) on FIELD_DEFINITION
+
 type Query {
     hero: Character @deprecated
-    droid(id: ID!): Droid
-    search(name: String!): SearchResult
+    droid(id: ID!): Droid @approx(tolerance: null)
+    search(name: String!): SearchResult @approx
     searchResults: [SearchResult]
 }
 
@@ -52,7 +56,7 @@ type Human implements Character {
 }
 
 type Droid implements Character {
-    name: String!
+    name: String! @approx
     primaryFunction: String!
     friends: [Character]
 }