Presence test cost tracking options

cel/env.go

@@ -462,12 +462,12 @@ func (e *Env) ResidualAst(a *Ast, details *EvalDetails) (*Ast, error) {
 
 // EstimateCost estimates the cost of a type checked CEL expression using the length estimates of input data and
 // extension functions provided by estimator.
-func (e *Env) EstimateCost(ast *Ast, estimator checker.CostEstimator) (checker.CostEstimate, error) {
+func (e *Env) EstimateCost(ast *Ast, estimator checker.CostEstimator, opts ...checker.CostOption) (checker.CostEstimate, error) {
 	checked, err := AstToCheckedExpr(ast)
 	if err != nil {
 		return checker.CostEstimate{}, fmt.Errorf("EsimateCost could not inspect Ast: %v", err)
 	}
-	return checker.Cost(checked, estimator), nil
+	return checker.Cost(checked, estimator, opts...)
 }
 
 // configure applies a series of EnvOptions to the current environment.

checker/cost.go

@@ -261,6 +261,8 @@ type coster struct {
 	computedSizes map[int64]SizeEstimate
 	checkedExpr   *exprpb.CheckedExpr
 	estimator     CostEstimator
+	// presenceTestCost will either be a zero or one based on whether has() macros count against cost computations.
+	presenceTestCost CostEstimate
 }
 
 // Use a stack of iterVar -> iterRange Expr Ids to handle shadowed variable names.
@@ -283,16 +285,39 @@ func (vs iterRangeScopes) peek(varName string) (int64, bool) {
 	return 0, false
 }
 
+// CostOption configures flags which affect cost computations.
+type CostOption func(*coster) error
+
+// PresenceTestHasCost determines whether presence testing has a cost of one or zero.
+// Defaults to presence test has a cost of one.
+func PresenceTestHasCost(hasCost bool) CostOption {
+	return func(c *coster) error {
+		if hasCost {
+			c.presenceTestCost = selectAndIdentCost
+			return nil
+		}
+		c.presenceTestCost = CostEstimate{Min: 0, Max: 0}
+		return nil
+	}
+}
+
 // Cost estimates the cost of the parsed and type checked CEL expression.
-func Cost(checker *exprpb.CheckedExpr, estimator CostEstimator) CostEstimate {
-	c := coster{
-		checkedExpr:   checker,
-		estimator:     estimator,
-		exprPath:      map[int64][]string{},
-		iterRanges:    map[string][]int64{},
-		computedSizes: map[int64]SizeEstimate{},
+func Cost(checker *exprpb.CheckedExpr, estimator CostEstimator, opts ...CostOption) (CostEstimate, error) {
+	c := &coster{
+		checkedExpr:      checker,
+		estimator:        estimator,
+		exprPath:         map[int64][]string{},
+		iterRanges:       map[string][]int64{},
+		computedSizes:    map[int64]SizeEstimate{},
+		presenceTestCost: CostEstimate{Min: 1, Max: 1},
+	}
+	for _, opt := range opts {
+		err := opt(c)
+		if err != nil {
+			return CostEstimate{}, err
+		}
 	}
-	return c.cost(checker.GetExpr())
+	return c.cost(checker.GetExpr()), nil
 }
 
 func (c *coster) cost(e *exprpb.Expr) CostEstimate {
@@ -347,7 +372,7 @@ func (c *coster) costSelect(e *exprpb.Expr) CostEstimate {
 		// this is equivalent to how evalTestOnly increments the runtime cost counter
 		// but does not add any additional cost for the qualifier, except here we do
 		// the reverse (ident adds cost)
-		sum = sum.Add(selectAndIdentCost)
+		sum = sum.Add(c.presenceTestCost)
 		sum = sum.Add(c.cost(sel.GetOperand()))
 		return sum
 	}

checker/cost_test.go

@@ -42,11 +42,12 @@ func TestCost(t *testing.T) {
 	zeroCost := CostEstimate{}
 	oneCost := CostEstimate{Min: 1, Max: 1}
 	cases := []struct {
-		name   string
-		expr   string
-		decls  []*exprpb.Decl
-		hints  map[string]int64
-		wanted CostEstimate
+		name    string
+		expr    string
+		decls   []*exprpb.Decl
+		hints   map[string]int64
+		options []CostOption
+		wanted  CostEstimate
 	}{
 		{
 			name:   "const",
@@ -71,12 +72,33 @@ func TestCost(t *testing.T) {
 			decls:  []*exprpb.Decl{decls.NewVar("input", allTypes)},
 			wanted: CostEstimate{Min: 2, Max: 2},
 		},
+		{
+			name:    "select: field test only no has() cost",
+			expr:    `has(input.single_int32)`,
+			decls:   []*exprpb.Decl{decls.NewVar("input", decls.NewObjectType("google.expr.proto3.test.TestAllTypes"))},
+			wanted:  CostEstimate{Min: 1, Max: 1},
+			options: []CostOption{PresenceTestHasCost(false)},
+		},
 		{
 			name:   "select: field test only",
 			expr:   `has(input.single_int32)`,
 			decls:  []*exprpb.Decl{decls.NewVar("input", decls.NewObjectType("google.expr.proto3.test.TestAllTypes"))},
 			wanted: CostEstimate{Min: 2, Max: 2},
 		},
+		{
+			name:    "select: non-proto field test has() cost",
+			expr:    `has(input.testAttr.nestedAttr)`,
+			decls:   []*exprpb.Decl{decls.NewVar("input", nestedMap)},
+			wanted:  CostEstimate{Min: 3, Max: 3},
+			options: []CostOption{PresenceTestHasCost(true)},
+		},
+		{
+			name:    "select: non-proto field test no has() cost",
+			expr:    `has(input.testAttr.nestedAttr)`,
+			decls:   []*exprpb.Decl{decls.NewVar("input", nestedMap)},
+			wanted:  CostEstimate{Min: 2, Max: 2},
+			options: []CostOption{PresenceTestHasCost(false)},
+		},
 		{
 			name:   "select: non-proto field test",
 			expr:   `has(input.testAttr.nestedAttr)`,
@@ -480,7 +502,10 @@ func TestCost(t *testing.T) {
 			if len(errs.GetErrors()) != 0 {
 				t.Fatalf("Check(%s) failed: %v", tc.expr, errs.ToDisplayString())
 			}
-			est := Cost(checked, testCostEstimator{hints: tc.hints})
+			est, err := Cost(checked, testCostEstimator{hints: tc.hints}, tc.options...)
+			if err != nil {
+				t.Fatalf("Cost() failed: %v", err)
+			}
 			if est.Min != tc.wanted.Min || est.Max != tc.wanted.Max {
 				t.Fatalf("Got cost interval [%v, %v], wanted [%v, %v]",
 					est.Min, est.Max, tc.wanted.Min, tc.wanted.Max)

interpreter/runtimecost.go

@@ -53,6 +53,11 @@ func CostObserver(tracker *CostTracker) EvalObserver {
 				tracker.stack.drop(t.Attr().ID())
 				tracker.cost += common.SelectAndIdentCost
 			}
+			if !tracker.presenceTestHasCost {
+				if _, isTestOnly := programStep.(*evalTestOnly); isTestOnly {
+					tracker.cost -= common.SelectAndIdentCost
+				}
+			}
 		case *evalExhaustiveConditional:
 			// Ternary has no direct cost. All cost is from the conditional and the true/false branch expressions.
 			tracker.stack.drop(t.attr.falsy.ID(), t.attr.truthy.ID(), t.attr.expr.ID())
@@ -95,21 +100,58 @@ func CostObserver(tracker *CostTracker) EvalObserver {
 	return observer
 }
 
-// CostTracker represents the information needed for tacking runtime cost
+// CostTrackerOption configures the behavior of CostTracker objects.
+type CostTrackerOption func(*CostTracker) error
+
+// CostTrackerLimit sets the runtime limit on the evaluation cost during execution and will terminate the expression
+// evaluation if the limit is exceeded.
+func CostTrackerLimit(limit uint64) CostTrackerOption {
+	return func(tracker *CostTracker) error {
+		tracker.Limit = &limit
+		return nil
+	}
+}
+
+// PresenceTestHasCost determines whether presence testing has a cost of one or zero.
+// Defaults to presence test has a cost of one.
+func PresenceTestHasCost(hasCost bool) CostTrackerOption {
+	return func(tracker *CostTracker) error {
+		tracker.presenceTestHasCost = hasCost
+		return nil
+	}
+}
+
+// NewCostTracker creates a new CostTracker with a given estimator and a set of functional CostTrackerOption values.
+func NewCostTracker(estimator ActualCostEstimator, opts ...CostTrackerOption) (*CostTracker, error) {
+	tracker := &CostTracker{
+		Estimator:           estimator,
+		presenceTestHasCost: true,
+	}
+	for _, opt := range opts {
+		err := opt(tracker)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return tracker, nil
+}
+
+// CostTracker represents the information needed for tracking runtime cost.
 type CostTracker struct {
-	Estimator ActualCostEstimator
-	Limit     *uint64
+	Estimator           ActualCostEstimator
+	Limit               *uint64
+	presenceTestHasCost bool
 
 	cost  uint64
 	stack refValStack
 }
 
 // ActualCost returns the runtime cost
-func (c CostTracker) ActualCost() uint64 {
+func (c *CostTracker) ActualCost() uint64 {
 	return c.cost
 }
 
-func (c CostTracker) costCall(call InterpretableCall, argValues []ref.Val, result ref.Val) uint64 {
+func (c *CostTracker) costCall(call InterpretableCall, argValues []ref.Val, result ref.Val) uint64 {
 	var cost uint64
 	if c.Estimator != nil {
 		callCost := c.Estimator.CallCost(call.Function(), call.OverloadID(), argValues, result)
@@ -179,7 +221,7 @@ func (c CostTracker) costCall(call InterpretableCall, argValues []ref.Val, resul
 }
 
 // actualSize returns the size of value
-func (c CostTracker) actualSize(value ref.Val) uint64 {
+func (c *CostTracker) actualSize(value ref.Val) uint64 {
 	if sz, ok := value.(traits.Sizer); ok {
 		return uint64(sz.Size().(types.Int))
 	}

interpreter/runtimecost_test.go

@@ -108,7 +108,7 @@ func TestTrackCostAdvanced(t *testing.T) {
 	}
 }
 
-func computeCost(t *testing.T, expr string, decls []*exprpb.Decl, ctx Activation, limit *uint64) (cost uint64, est checker.CostEstimate, err error) {
+func computeCost(t *testing.T, expr string, decls []*exprpb.Decl, ctx Activation, options []CostTrackerOption) (cost uint64, est checker.CostEstimate, err error) {
 	t.Helper()
 
 	s := common.NewTextSource(expr)
@@ -129,14 +129,19 @@ func computeCost(t *testing.T, expr string, decls []*exprpb.Decl, ctx Activation
 	if err != nil {
 		t.Fatalf("Failed to initialize env: %v", err)
 	}
-
+	costTracker, err := NewCostTracker(&testRuntimeCostEstimator{}, options...)
+	if err != nil {
+		t.Fatalf("NewCostTracker() failed: %v", err)
+	}
 	checked, errs := checker.Check(parsed, s, env)
 	if len(errs.GetErrors()) != 0 {
 		t.Fatalf(`Failed to check expression "%s", error: %v`, expr, errs.GetErrors())
 	}
-	est = checker.Cost(checked, testCostEstimator{})
+	est, err = checker.Cost(checked, testCostEstimator{}, checker.PresenceTestHasCost(costTracker.presenceTestHasCost))
+	if err != nil {
+		t.Fatalf("checker.Cost() failed: %v", err)
+	}
 	interp := NewStandardInterpreter(cont, reg, reg, attrs)
-	costTracker := &CostTracker{Estimator: &testRuntimeCostEstimator{}, Limit: limit}
 	prg, err := interp.NewInterpretable(checked, Observe(CostObserver(costTracker)))
 	if err != nil {
 		t.Fatalf(`Failed to check expression "%s", error: %v`, expr, errs.GetErrors())
@@ -244,6 +249,7 @@ func TestRuntimeCost(t *testing.T) {
 		in           any
 		testFuncCost bool
 		limit        uint64
+		options      []CostTrackerOption
 
 		expectExceedsLimit bool
 	}{
@@ -302,6 +308,22 @@ func TestRuntimeCost(t *testing.T) {
 			want:  3,
 			in:    map[string]any{"input": []string{"v"}},
 		},
+		{
+			name:    "select: field test only no has() cost",
+			expr:    `has(input.single_int32)`,
+			decls:   []*exprpb.Decl{decls.NewVar("input", decls.NewObjectType("google.expr.proto3.test.TestAllTypes"))},
+			want:    1,
+			options: []CostTrackerOption{PresenceTestHasCost(false)},
+			in: map[string]any{
+				"input": &proto3pb.TestAllTypes{
+					RepeatedBool: []bool{false},
+					MapInt64NestedType: map[int64]*proto3pb.NestedTestAllTypes{
+						1: {},
+					},
+					MapStringString: map[string]string{},
+				},
+			},
+		},
 		{
 			name:  "select: field test only",
 			expr:  `has(input.single_int32)`,
@@ -317,6 +339,34 @@ func TestRuntimeCost(t *testing.T) {
 				},
 			},
 		},
+		{
+			name:    "select: non-proto field test has() cost",
+			expr:    `has(input.testAttr.nestedAttr)`,
+			decls:   []*exprpb.Decl{decls.NewVar("input", nestedMap)},
+			want:    3,
+			options: []CostTrackerOption{PresenceTestHasCost(true)},
+			in: map[string]any{
+				"input": map[string]any{
+					"testAttr": map[string]any{
+						"nestedAttr": "0",
+					},
+				},
+			},
+		},
+		{
+			name:    "select: non-proto field test no has() cost",
+			expr:    `has(input.testAttr.nestedAttr)`,
+			decls:   []*exprpb.Decl{decls.NewVar("input", nestedMap)},
+			want:    2,
+			options: []CostTrackerOption{PresenceTestHasCost(false)},
+			in: map[string]any{
+				"input": map[string]any{
+					"testAttr": map[string]any{
+						"nestedAttr": "0",
+					},
+				},
+			},
+		},
 		{
 			name:  "select: non-proto field test",
 			expr:  `has(input.testAttr.nestedAttr)`,
@@ -729,7 +779,11 @@ func TestRuntimeCost(t *testing.T) {
 			if tc.limit > 0 {
 				costLimit = &tc.limit
 			}
-			actualCost, est, err := computeCost(t, tc.expr, tc.decls, ctx, costLimit)
+			options := tc.options
+			if costLimit != nil {
+				options = append(options, CostTrackerLimit(*costLimit))
+			}
+			actualCost, est, err := computeCost(t, tc.expr, tc.decls, ctx, options)
 			if err != nil {
 				if tc.expectExceedsLimit {
 					return