From 04f6a526ea2848cce7f215ee3dc20cba5fe704bf Mon Sep 17 00:00:00 2001 From: TristonianJones Date: Thu, 27 Aug 2026 13:19:26 -0700 Subject: [PATCH 1/6] Standard operations in the new unified cost model --- common/cost/standard.go | 169 +++++++++++++++++ common/cost/standard_test.go | 339 +++++++++++++++++++++++++++++++++++ 2 files changed, 508 insertions(+) create mode 100644 common/cost/standard.go create mode 100644 common/cost/standard_test.go diff --git a/common/cost/standard.go b/common/cost/standard.go new file mode 100644 index 00000000..d6bc6581 --- /dev/null +++ b/common/cost/standard.go @@ -0,0 +1,169 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package cost + +import "cel.dev/cel-go/common/overloads" + +// StandardOverloadModels defines the cost models for standard CEL functions. +var StandardOverloadModels = []OverloadModel{ + // O(1) container index operations + Overload(overloads.IndexList, + EvalCost(Const(1)), + ResultSize(ArgElem(0)), + ), + Overload(overloads.IndexMap, + EvalCost(Const(1)), + ResultSize(ArgElem(0)), + ), + + // O(n) prefix/suffix functions + MemberOverload(overloads.StartsWithString, + EvalCost(Scale(Arg(0), StringTraversalCostFactor)), + ), + MemberOverload(overloads.EndsWithString, + EvalCost(Scale(Arg(0), StringTraversalCostFactor)), + ), + + // O(n) conversion & format functions + Overload(overloads.StringToBytes, + EvalCost(Scale(Arg(0), StringTraversalCostFactor)), + ResultSize(Ranged(Arg(0), Scale(Arg(0), 4.0))), + ), + Overload(overloads.BytesToString, + EvalCost(Scale(Arg(0), StringTraversalCostFactor)), + ResultSize(Ranged(Scale(Arg(0), 0.25), Arg(0))), + ), + Overload(overloads.ExtQuoteString, + EvalCost(Scale(Arg(0), StringTraversalCostFactor)), + ResultSize(Ranged(Sum(Arg(0), Const(2)), Sum(Scale(Arg(0), 2.0), Const(2)))), + ), + MemberOverload(overloads.ExtFormatString, + EvalCost(Scale(Target(), StringTraversalCostFactor)), + ), + + // O(n) containment + Overload(overloads.InList, EvalCost(Arg(1))), + + // O(min(m, n)) comparison / equality + Overload(overloads.LessString, + EvalCost(Scale(Min(Arg(0), Arg(1)), StringTraversalCostFactor)), + ), + Overload(overloads.GreaterString, + EvalCost(Scale(Min(Arg(0), Arg(1)), StringTraversalCostFactor)), + ), + Overload(overloads.LessEqualsString, + EvalCost(Scale(Min(Arg(0), Arg(1)), StringTraversalCostFactor)), + ), + Overload(overloads.GreaterEqualsString, + EvalCost(Scale(Min(Arg(0), Arg(1)), StringTraversalCostFactor)), + ), + Overload(overloads.LessBytes, + EvalCost(Scale(Min(Arg(0), Arg(1)), StringTraversalCostFactor)), + ), + Overload(overloads.GreaterBytes, + EvalCost(Scale(Min(Arg(0), Arg(1)), StringTraversalCostFactor)), + ), + Overload(overloads.LessEqualsBytes, + EvalCost(Scale(Min(Arg(0), Arg(1)), StringTraversalCostFactor)), + ), + Overload(overloads.GreaterEqualsBytes, + EvalCost(Scale(Min(Arg(0), Arg(1)), StringTraversalCostFactor)), + ), + Overload(overloads.Equals, + EvalCost(Scale(Min(Arg(0), Arg(1)), StringTraversalCostFactor)), + ), + Overload(overloads.NotEquals, + EvalCost(Scale(Min(Arg(0), Arg(1)), StringTraversalCostFactor)), + ), + + // O(m+n) string & bytes concatenation + Overload(overloads.AddString, + EvalCost(Scale(Sum(Arg(0), Arg(1)), StringTraversalCostFactor)), + ResultSize(Sum(Arg(0), Arg(1))), + ), + Overload(overloads.AddBytes, + EvalCost(Scale(Sum(Arg(0), Arg(1)), StringTraversalCostFactor)), + ResultSize(Sum(Arg(0), Arg(1))), + ), + + // O(1) list concatenation with size tracking + Overload(overloads.AddList, + EvalCost(Const(1)), + ResultSize(Sum(Arg(0), Arg(1))), + ), + + // O(nm) regex matches + Overload(overloads.Matches, + EvalCost(Mul( + Scale(Sum(Arg(0), Const(1)), StringTraversalCostFactor), + Scale(Arg(1), RegexStringLengthCostFactor), + )), + ), + MemberOverload(overloads.MatchesString, + EvalCost(Mul( + Scale(Sum(Target(), Const(1)), StringTraversalCostFactor), + Scale(Arg(0), RegexStringLengthCostFactor), + )), + ), + + // O(nm) substring contains + MemberOverload(overloads.ContainsString, + EvalCost(Mul( + Scale(Target(), StringTraversalCostFactor), + Scale(Arg(0), StringTraversalCostFactor), + )), + ), + + // The arg cost for logical and conditional operators is special-cased for + // short-circuiting (see CalculateArgCost in estimator.go), so the arg cost is 0. + + // Logical short-circuiting operators + Overload(overloads.LogicalOr, EvalCost(Const(0))), + Overload(overloads.LogicalAnd, EvalCost(Const(0))), + + // Conditional operator + Overload(overloads.Conditional, + EvalCost(Const(0)), + ResultSize(Union(Arg(1), Arg(2))), + ), +} + +// StandardOverloadEstimators returns the map of FunctionEstimator instances for standard overloads. +func StandardOverloadEstimators() map[string]FunctionEstimator { + return StandardOverloadEstimatorsWithOptions(nil) +} + +// StandardOverloadEstimatorsWithOptions returns the map of FunctionEstimator instances for standard overloads with an optional SizingStrategy. +func StandardOverloadEstimatorsWithOptions(strategy SizingStrategy) map[string]FunctionEstimator { + estimators := make(map[string]FunctionEstimator, len(StandardOverloadModels)) + for _, m := range StandardOverloadModels { + estimators[m.ID] = m.FunctionEstimatorWithOptions(strategy) + } + return estimators +} + +// StandardOverloadTrackers returns the map of FunctionTracker instances for standard overloads. +func StandardOverloadTrackers() map[string]FunctionTracker { + return StandardOverloadTrackersWithOptions(nil) +} + +// StandardOverloadTrackersWithOptions returns the map of FunctionTracker instances for standard overloads with an optional SizingStrategy. +func StandardOverloadTrackersWithOptions(strategy SizingStrategy) map[string]FunctionTracker { + trackers := make(map[string]FunctionTracker, len(StandardOverloadModels)) + for _, m := range StandardOverloadModels { + trackers[m.ID] = m.FunctionTrackerWithOptions(strategy) + } + return trackers +} diff --git a/common/cost/standard_test.go b/common/cost/standard_test.go new file mode 100644 index 00000000..1a59b47c --- /dev/null +++ b/common/cost/standard_test.go @@ -0,0 +1,339 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package cost + +import ( + "testing" + + "cel.dev/cel-go/common/ast" + "cel.dev/cel-go/common/overloads" + "cel.dev/cel-go/common/types" + "cel.dev/cel-go/common/types/ref" +) + +func TestStandardOverloadModels_CollectionCompleteness(t *testing.T) { + estimators := StandardOverloadEstimators() + trackers := StandardOverloadTrackers() + + if len(estimators) != len(StandardOverloadModels) { + t.Errorf("got %d estimators, wanted %d", len(estimators), len(StandardOverloadModels)) + } + if len(trackers) != len(StandardOverloadModels) { + t.Errorf("got %d trackers, wanted %d", len(trackers), len(StandardOverloadModels)) + } +} + +func TestStandardOverloadModels_BasicOperationTrackers(t *testing.T) { + trackers := StandardOverloadTrackers() + adapter := types.DefaultTypeAdapter + + tests := []struct { + name string + overloadID string + args []ref.Val + wantCost uint64 + }{ + { + name: "in_list_string_elements", + overloadID: overloads.InList, + args: []ref.Val{ + types.String("item"), + adapter.NativeToValue([]string{"a", "b", "c"}), + }, + wantCost: 3, + }, + { + name: "index_list_lookup", + overloadID: overloads.IndexList, + args: []ref.Val{ + adapter.NativeToValue([]string{"a", "b", "c"}), + types.Int(1), + }, + wantCost: 1, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + tracker := trackers[tc.overloadID] + if tracker == nil { + t.Fatalf("missing tracker for %s", tc.overloadID) + } + cost := tracker(tc.args, nil) + if cost == nil { + t.Fatalf("tracker returned nil cost") + } + if *cost != tc.wantCost { + t.Errorf("cost = %d, want %d", *cost, tc.wantCost) + } + }) + } +} + +func TestEqualsNotEqualsOverloadModels_Trackers(t *testing.T) { + trackers := StandardOverloadTrackers() + eqTracker := trackers[overloads.Equals] + if eqTracker == nil { + t.Fatalf("missing tracker for Equals") + } + neTracker := trackers[overloads.NotEquals] + if neTracker == nil { + t.Fatalf("missing tracker for NotEquals") + } + + adapter := types.DefaultTypeAdapter + + largeIntSlice := make([]int64, 40) + for i := range largeIntSlice { + largeIntSlice[i] = int64(i) + } + + largeByteSlice := make([]byte, 50) + for i := range largeByteSlice { + largeByteSlice[i] = byte(i) + } + + tests := []struct { + name string + tracker FunctionTracker + lhs, rhs any + wantCost uint64 + }{ + { + name: "int_list_equal_small", + tracker: eqTracker, + lhs: []int64{1, 2, 3}, + rhs: []int64{1, 2, 3}, + wantCost: 1, // ceil(3 * 0.1) = 1 + }, + { + name: "int_list_equal_large", + tracker: eqTracker, + lhs: largeIntSlice, + rhs: largeIntSlice, + wantCost: 4, // ceil(40 * 0.1) = 4 + }, + { + name: "int_list_unequal_sizes", + tracker: neTracker, + lhs: []int64{1, 2, 3, 4, 5}, + rhs: largeIntSlice, + wantCost: 1, // ceil(min(5, 40) * 0.1) = 1 + }, + { + name: "map_equal", + tracker: eqTracker, + lhs: map[string]int64{"a": 1, "b": 2}, + rhs: map[string]int64{"a": 1, "b": 2}, + wantCost: 1, // ceil(2 * 0.1) = 1 + }, + { + name: "bytes_equal_small", + tracker: eqTracker, + lhs: []byte("hello"), + rhs: []byte("hello"), + wantCost: 1, // ceil(5 * 0.1) = 1 + }, + { + name: "bytes_not_equal_large", + tracker: neTracker, + lhs: largeByteSlice, + rhs: largeByteSlice, + wantCost: 5, // ceil(50 * 0.1) = 5 + }, + { + name: "string_equal", + tracker: eqTracker, + lhs: "hello world", + rhs: "hello world", + wantCost: 2, // ceil(11 * 0.1) = 2 + }, + { + name: "scalar_int_equal", + tracker: eqTracker, + lhs: int64(42), + rhs: int64(42), + wantCost: 1, // ceil(1 * 0.1) = 1 + }, + { + name: "scalar_bool_not_equal", + tracker: neTracker, + lhs: true, + rhs: false, + wantCost: 1, // ceil(1 * 0.1) = 1 + }, + { + name: "scalar_double_equal", + tracker: eqTracker, + lhs: 3.14, + rhs: 3.14, + wantCost: 1, // ceil(1 * 0.1) = 1 + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + args := []ref.Val{ + adapter.NativeToValue(tc.lhs), + adapter.NativeToValue(tc.rhs), + } + cost := tc.tracker(args, nil) + if cost == nil { + t.Errorf("cost got nil, wanted %d", tc.wantCost) + } else if *cost != tc.wantCost { + t.Errorf("cost got %d, wanted %d", *cost, tc.wantCost) + } + }) + } +} + +func TestEqualsNotEqualsOverloadModels_Estimators(t *testing.T) { + estimators := StandardOverloadEstimators() + eqEstimator := estimators[overloads.Equals] + if eqEstimator == nil { + t.Fatalf("missing estimator for Equals") + } + neEstimator := estimators[overloads.NotEquals] + if neEstimator == nil { + t.Fatalf("missing estimator for NotEquals") + } + + tests := []struct { + name string + estimator FunctionEstimator + nodeA AstNode + nodeB AstNode + wantMin uint64 + wantMax uint64 + }{ + { + name: "int_list_nodes_equal", + estimator: eqEstimator, + nodeA: &testAstNode{t: types.NewListType(types.IntType), size: &SizeEstimate{Min: 10, Max: 40}}, + nodeB: &testAstNode{t: types.NewListType(types.IntType), size: &SizeEstimate{Min: 20, Max: 30}}, + wantMin: 1, // ceil(min(10, 20) * 0.1) = 1 + wantMax: 3, // ceil(min(40, 30) * 0.1) = 3 + }, + { + name: "map_nodes_not_equal", + estimator: neEstimator, + nodeA: &testAstNode{t: types.NewMapType(types.StringType, types.IntType), size: &SizeEstimate{Min: 15, Max: 50}}, + nodeB: &testAstNode{t: types.NewMapType(types.StringType, types.IntType), size: &SizeEstimate{Min: 5, Max: 60}}, + wantMin: 1, // ceil(min(15, 5) * 0.1) = 1 + wantMax: 5, // ceil(min(50, 60) * 0.1) = 5 + }, + { + name: "scalar_int_nodes_equal", + estimator: eqEstimator, + nodeA: &testAstNode{t: types.IntType, size: &SizeEstimate{Min: 1, Max: 1}}, + nodeB: &testAstNode{t: types.IntType, size: &SizeEstimate{Min: 1, Max: 1}}, + wantMin: 1, + wantMax: 1, + }, + { + name: "string_nodes_equal", + estimator: eqEstimator, + nodeA: &testAstNode{t: types.StringType, size: &SizeEstimate{Min: 10, Max: 40}}, + nodeB: &testAstNode{t: types.StringType, size: &SizeEstimate{Min: 20, Max: 30}}, + wantMin: 1, // ceil(min(10, 20) * 0.1) = 1 + wantMax: 3, // ceil(min(40, 30) * 0.1) = 3 + }, + { + name: "bytes_nodes_equal", + estimator: eqEstimator, + nodeA: &testAstNode{t: types.BytesType, size: &SizeEstimate{Min: 20, Max: 80}}, + nodeB: &testAstNode{t: types.BytesType, size: &SizeEstimate{Min: 10, Max: 100}}, + wantMin: 1, // ceil(min(20, 10) * 0.1) = 1 + wantMax: 8, // ceil(min(80, 100) * 0.1) = 8 + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + res := tc.estimator(nil, nil, []AstNode{tc.nodeA, tc.nodeB}) + if res == nil { + t.Fatalf("estimator returned nil") + } + if res.CostEstimate.Min != tc.wantMin || res.CostEstimate.Max != tc.wantMax { + t.Errorf("estimator cost got [%d, %d], wanted [%d, %d]", + res.CostEstimate.Min, res.CostEstimate.Max, tc.wantMin, tc.wantMax) + } + }) + } +} + +type testAstNode struct { + path []string + t *types.Type + size *SizeEstimate +} + +func (n *testAstNode) Path() []string { return n.path } +func (n *testAstNode) Type() *types.Type { return n.t } +func (n *testAstNode) Expr() ast.Expr { return nil } +func (n *testAstNode) ComputedSize() *SizeEstimate { return n.size } + +func TestOverloadConstructors(t *testing.T) { + tests := []struct { + name string + model OverloadModel + wantID string + wantMember bool + wantTarget bool + }{ + { + name: "custom_global_overload", + model: Overload("custom_global", + EvalCost(Scale(Arg(0), 1.5)), + ResultSize(Sum(Arg(0), Const(1))), + ), + wantID: "custom_global", + wantMember: false, + wantTarget: false, + }, + { + name: "custom_member_overload", + model: MemberOverload("custom_member", + EvalCost(Scale(Arg(0), 2.0)), + ), + wantID: "custom_member", + wantMember: true, + wantTarget: true, + }, + { + name: "inferred_target_overload", + model: Overload("inferred_member", + EvalCost(Scale(Target(), 2.0)), + ), + wantID: "inferred_member", + wantMember: false, + wantTarget: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if tc.model.ID != tc.wantID { + t.Errorf("ID = %q, want %q", tc.model.ID, tc.wantID) + } + if tc.model.IsMember != tc.wantMember { + t.Errorf("IsMember = %t, want %t", tc.model.IsMember, tc.wantMember) + } + if tc.model.hasTarget() != tc.wantTarget { + t.Errorf("hasTarget() = %t, want %t", tc.model.hasTarget(), tc.wantTarget) + } + }) + } +} From 365afa48552c8c01831a8c8592cbd7598587b661 Mon Sep 17 00:00:00 2001 From: TristonianJones Date: Thu, 27 Aug 2026 13:28:18 -0700 Subject: [PATCH 2/6] Update cost estimation and tracking to use standard overload models --- common/cost/BUILD.bazel | 2 + common/cost/estimator.go | 652 +++++++++++++++------------------- common/cost/estimator_test.go | 519 ++++++++++++++++++++++++++- common/cost/tracker.go | 123 +++---- common/cost/tracker_test.go | 322 ++++++++++++++--- 5 files changed, 1129 insertions(+), 489 deletions(-) diff --git a/common/cost/BUILD.bazel b/common/cost/BUILD.bazel index ba078cef..38cbbcf5 100644 --- a/common/cost/BUILD.bazel +++ b/common/cost/BUILD.bazel @@ -13,6 +13,7 @@ go_library( "default_strategy.go", "estimator.go", "model.go", + "standard.go", "strategy.go", "tracker.go", ], @@ -36,6 +37,7 @@ go_test( "default_strategy_test.go", "estimator_test.go", "model_test.go", + "standard_test.go", "tracker_test.go", ], embed = [ diff --git a/common/cost/estimator.go b/common/cost/estimator.go index 82da6b9f..0f4ad355 100644 --- a/common/cost/estimator.go +++ b/common/cost/estimator.go @@ -18,32 +18,27 @@ import ( "math" "cel.dev/cel-go/common/ast" + "cel.dev/cel-go/common/operators" "cel.dev/cel-go/common/overloads" "cel.dev/cel-go/common/types" ) -// WARNING: Any changes to cost calculations in this file require a corresponding change in interpreter/runtimecost.go - -// Estimator estimates the sizes of variable length input data and the costs of functions. +// Estimator provides CallCost and Size estimates for cost calculation. type Estimator interface { - // EstimateSize returns a SizeEstimate for the given AstNode, or nil if the estimator has no - // estimate to provide. + + // EstimateCallCost returns the CallEstimate for a function, overload and target and arguments. // - // The size is equivalent to the result of the CEL `size()` function: - // * Number of unicode characters in a string - // * Number of bytes in a sequence - // * Number of map entries or number of list items. + // If nil is returned, the cost of the function is calculated by the OverloadCostEstimate if + // provided, or standard CEL function cost estimation. + EstimateCallCost(function, overloadID string, target *AstNode, args []AstNode) *CallEstimate + + // EstimateSize returns the SizeEstimate of an AstNode. // - // EstimateSize is only called for AstNodes where CEL does not know the size; EstimateSize is not - // called for values defined inline in CEL where the size is already obvious to CEL. + // If nil is returned, the size of the node is calculated by the SizingStrategy. EstimateSize(element AstNode) *SizeEstimate - - // EstimateCallCost returns the estimated cost of an invocation, or nil if the estimator has no - // estimate to provide. - EstimateCallCost(function, overloadID string, target *AstNode, args []AstNode) *CallEstimate } -// AstNode represents an AST node for the purpose of cost estimations. +// AstNode represents an AST node for estimating cost. type AstNode interface { // Path returns a field path through the provided type declarations to the type of the AstNode, or nil if the AstNode does not // represent type directly reachable from the provided type declarations. @@ -91,9 +86,9 @@ func (e astNode) ComputedSize() *SizeEstimate { return e.derivedSize } -// NewAstNode creates a new AstNode for cost estimation. +// NewAstNode returns an AstNode with the given expression, path, type, and derived size. func NewAstNode(expr ast.Expr, path []string, t *types.Type, derivedSize *SizeEstimate) AstNode { - return &astNode{ + return astNode{ path: path, t: t, expr: expr, @@ -101,13 +96,18 @@ func NewAstNode(expr ast.Expr, path []string, t *types.Type, derivedSize *SizeEs } } -// CostOption configures flags which affect cost computations. -type CostOption func(*coster) error +// Option configures flags which affect cost computations. +type Option func(*coster) error + +// CostOption is an alias for Option. +// +// Deprecated: use Option. +type CostOption = Option // 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 { +func PresenceTestHasCost(hasCost bool) Option { return func(c *coster) error { if hasCost { c.presenceTestCost = selectAndIdentCost @@ -125,15 +125,23 @@ type FunctionEstimator func(estimator Estimator, target *AstNode, args []AstNode // // When a OverloadCostEstimate is provided, it will override the cost calculation of the CostEstimator provided to // the Cost() call. -func OverloadCostEstimate(overloadID string, functionCoster FunctionEstimator) CostOption { +func OverloadCostEstimate(overloadID string, functionCoster FunctionEstimator) Option { return func(c *coster) error { c.overloadEstimators[overloadID] = functionCoster return nil } } +// EstimateSizingStrategy configures a custom SizingStrategy for cost estimation. +func EstimateSizingStrategy(strategy SizingStrategy) Option { + return func(c *coster) error { + c.sizingStrategy = strategy + return nil + } +} + // Cost estimates the cost of the parsed and type checked CEL expression. -func Cost(checked *ast.AST, estimator Estimator, opts ...CostOption) (CostEstimate, error) { +func Cost(checked *ast.AST, estimator Estimator, opts ...Option) (CostEstimate, error) { c := &coster{ checkedAST: checked, estimator: estimator, @@ -141,7 +149,6 @@ func Cost(checked *ast.AST, estimator Estimator, opts ...CostOption) (CostEstima exprPaths: map[int64][]string{}, localVars: make(scopes), computedSizes: map[int64]SizeEstimate{}, - computedEntrySizes: map[int64]entrySizeEstimate{}, presenceTestCost: FixedCostEstimate(1), } for _, opt := range opts { @@ -161,82 +168,32 @@ type coster struct { localVars scopes // computedSizes tracks the computed sizes of call results. computedSizes map[int64]SizeEstimate - // computedEntrySizes tracks the size of list and map entries - computedEntrySizes map[int64]entrySizeEstimate - checkedAST *ast.AST - estimator Estimator - overloadEstimators map[string]FunctionEstimator + checkedAST *ast.AST + estimator Estimator + sizingStrategy SizingStrategy + overloadEstimators map[string]FunctionEstimator + sizingOverloadEstimators map[string]FunctionEstimator // presenceTestCost will either be a zero or one based on whether has() macros count against cost computations. presenceTestCost CostEstimate } -// entrySizeEstimate captures the container kind and associated key/index and value SizeEstimate values. -// -// An entrySizeEstimate only exists if both the key/index and the value have SizeEstimate values, otherwise -// a nil entrySizeEstimate should be used. -type entrySizeEstimate struct { - containerKind types.Kind - key SizeEstimate - val SizeEstimate -} - -// container returns the container kind (list or map) of the entry. -func (s *entrySizeEstimate) container() types.Kind { - if s == nil { - return types.UnknownKind - } - return s.containerKind -} - -// keySize returns the SizeEstimate for the key if one exists. -func (s *entrySizeEstimate) keySize() *SizeEstimate { - if s == nil { - return nil - } - return &s.key -} - -// valSize returns the SizeEstimate for the value if one exists. -func (s *entrySizeEstimate) valSize() *SizeEstimate { - if s == nil { - return nil - } - return &s.val -} - -// union returns the union of two entrySizeEstimates. -func (s *entrySizeEstimate) union(other *entrySizeEstimate) *entrySizeEstimate { - if s == nil || other == nil { - return nil - } - sk := s.key.Union(other.key) - sv := s.val.Union(other.val) - return &entrySizeEstimate{ - containerKind: s.containerKind, - key: sk, - val: sv, - } -} - -// localVar captures the local variable size and entrySize estimates if they exist for variables +// localVar captures the local variable size estimates if they exist for variables type localVar struct { - exprID int64 - path []string - size *SizeEstimate - entrySize *entrySizeEstimate + exprID int64 + path []string + size *SizeEstimate } // scopes is a stack of variable name to integer id stack to handle scopes created by cel.bind() like macros type scopes map[string][]*localVar -// push adds a variable name to the scope stack with its path, size, and entrySize estimates. -func (s scopes) push(varName string, expr ast.Expr, path []string, size *SizeEstimate, entrySize *entrySizeEstimate) { +// push adds a variable name to the scope stack with its path and size estimates. +func (s scopes) push(varName string, expr ast.Expr, path []string, size *SizeEstimate) { s[varName] = append(s[varName], &localVar{ - exprID: expr.ID(), - path: path, - size: size, - entrySize: entrySize, + exprID: expr.ID(), + path: path, + size: size, }) } @@ -255,58 +212,66 @@ func (s scopes) peek(varName string) (*localVar, bool) { return nil, false } -// containerKind returns the deduced container kind for a range expression. -func (c *coster) containerKind(rangeExpr ast.Expr, entrySize *entrySizeEstimate) types.Kind { - if k := entrySize.container(); k != types.UnknownKind { - return k - } - return c.getType(rangeExpr).Kind() -} - // pushIterKey pushes the iteration key or index variable for a comprehension onto the scope stack. func (c *coster) pushIterKey(varName string, rangeExpr ast.Expr) { - entrySize := c.computeEntrySize(rangeExpr) - size := entrySize.keySize() + rangeSize := c.sizeOrUnknown(rangeExpr) + var size *SizeEstimate + if rangeSize.Key != nil { + size = rangeSize.Key + } else { + s := FixedSizeEstimate(1) + size = &s + } path := c.getPath(rangeExpr) + container := c.getType(rangeExpr).Kind() subpath := "@keys" - if c.containerKind(rangeExpr, entrySize) == types.ListKind { + if container == types.ListKind { subpath = "@indices" } - c.localVars.push(varName, rangeExpr, append(path, subpath), size, nil) + c.localVars.push(varName, rangeExpr, append(path, subpath), size) } // pushIterValue pushes the iteration value variable for a comprehension onto the scope stack. func (c *coster) pushIterValue(varName string, rangeExpr ast.Expr) { - entrySize := c.computeEntrySize(rangeExpr) - size := entrySize.valSize() + rangeSize := c.sizeOrUnknown(rangeExpr) + var size *SizeEstimate + if rangeSize.Elem != nil { + size = rangeSize.Elem + } path := c.getPath(rangeExpr) + container := c.getType(rangeExpr).Kind() subpath := "@values" - if c.containerKind(rangeExpr, entrySize) == types.ListKind { + if container == types.ListKind { subpath = "@items" } - c.localVars.push(varName, rangeExpr, append(path, subpath), size, nil) + c.localVars.push(varName, rangeExpr, append(path, subpath), size) } // pushIterSingle pushes a single iteration variable (items for list, keys for map) onto the scope stack. func (c *coster) pushIterSingle(varName string, rangeExpr ast.Expr) { - entrySize := c.computeEntrySize(rangeExpr) - size := entrySize.keySize() + rangeSize := c.sizeOrUnknown(rangeExpr) + var size *SizeEstimate subpath := "@keys" - if c.containerKind(rangeExpr, entrySize) == types.ListKind { - size = entrySize.valSize() + container := c.getType(rangeExpr).Kind() + if container == types.ListKind { + size = rangeSize.Elem subpath = "@items" + } else { + if rangeSize.Key != nil { + size = rangeSize.Key + } else { + s := FixedSizeEstimate(1) + size = &s + } } path := c.getPath(rangeExpr) - c.localVars.push(varName, rangeExpr, append(path, subpath), size, nil) + c.localVars.push(varName, rangeExpr, append(path, subpath), size) } -// pushLocalVar records a local variable binding with its path, size, and entry size estimates. +// pushLocalVar records a local variable binding with its path and size estimates. func (c *coster) pushLocalVar(varName string, e ast.Expr) { path := c.getPath(e) - // note: retrieve the entry size for the local variable based on the size of the binding expression - // since the binding expression could be a list or map, the entry size should also be propagated - entrySize := c.computeEntrySize(e) - c.localVars.push(varName, e, path, c.computeSize(e), entrySize) + c.localVars.push(varName, e, path, c.computeSize(e)) } // peekLocalVar looks up the top of the scope stack for a local variable name. @@ -324,42 +289,38 @@ func (c *coster) cost(e ast.Expr) CostEstimate { if e == nil { return CostEstimate{} } - var estimate CostEstimate switch e.Kind() { case ast.LiteralKind: - estimate = constCost + return constCost case ast.IdentKind: - estimate = c.costIdent(e) + return c.costIdent(e) case ast.SelectKind: - estimate = c.costSelect(e) + return c.costSelect(e) case ast.CallKind: - estimate = c.costCall(e) + return c.costCall(e) case ast.ListKind: - estimate = c.costCreateList(e) + return c.costCreateList(e) case ast.MapKind: - estimate = c.costCreateMap(e) + return c.costCreateMap(e) case ast.StructKind: - estimate = c.costCreateStruct(e) + return c.costCreateStruct(e) case ast.ComprehensionKind: if c.isBind(e) { - estimate = c.costBind(e) - } else { - estimate = c.costComprehension(e) + return c.costBind(e) } + return c.costComprehension(e) default: return CostEstimate{} } - return estimate } // costIdent estimates the cost of evaluating an identifier and tracks its path. func (c *coster) costIdent(e ast.Expr) CostEstimate { - identName := e.AsIdent() - // build and track the field path - if v, ok := c.peekLocalVar(identName); ok { + ident := e.AsIdent() + if v, ok := c.peekLocalVar(ident); ok { c.addPath(e, v.path) } else { - c.addPath(e, []string{identName}) + c.addPath(e, []string{ident}) } return selectAndIdentCost } @@ -369,26 +330,43 @@ func (c *coster) costSelect(e ast.Expr) CostEstimate { sel := e.AsSelect() var sum CostEstimate if sel.IsTestOnly() { - // recurse, but do not add any cost - // 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(c.presenceTestCost) - sum = sum.Add(c.cost(sel.Operand())) - return sum + } else { + sum = sum.Add(selectAndIdentCost) } sum = sum.Add(c.cost(sel.Operand())) - targetType := c.getType(sel.Operand()) - switch targetType.Kind() { - case types.MapKind, types.StructKind, types.TypeParamKind: - sum = sum.Add(selectAndIdentCost) + sum = sum.Add(c.relativeAttributeCost(sel.Operand())) + targetPath := c.getPath(sel.Operand()) + if len(targetPath) > 0 { + c.addPath(e, append(targetPath, sel.FieldName())) } - - // build and track the field path - c.addPath(e, append(c.getPath(sel.Operand()), sel.FieldName())) return sum } +// relativeAttributeCost is the cost of qualifying a value which was computed rather than named. +func (c *coster) relativeAttributeCost(operand ast.Expr) CostEstimate { + if isAttributeChain(operand) { + return CostEstimate{} + } + return FixedCostEstimate(SelectAndIdentCost) +} + +// isAttributeChain reports whether an expression is resolved as part of a single attribute +// during evaluation. A chain begins at an identifier, or at a ternary which selects between +// attributes, and is extended by field selections and index operations. +func isAttributeChain(e ast.Expr) bool { + switch e.Kind() { + case ast.IdentKind, ast.SelectKind: + return true + case ast.CallKind: + switch e.AsCall().FunctionName() { + case operators.Index, operators.Conditional: + return true + } + } + return false +} + // costCall estimates the cost of evaluating a function call expression. func (c *coster) costCall(e ast.Expr) CostEstimate { // Dyn is just a way to disable type-checking, so return the cost of 1 with the cost of the argument @@ -401,6 +379,10 @@ func (c *coster) costCall(e ast.Expr) CostEstimate { args := call.Args() var sum CostEstimate + if call.FunctionName() == operators.Index && len(args) > 0 { + sum = sum.Add(c.relativeAttributeCost(args[0])) + } + argTypes := make([]AstNode, len(args)) argCosts := make([]CostEstimate, len(args)) for i, arg := range args { @@ -410,7 +392,11 @@ func (c *coster) costCall(e ast.Expr) CostEstimate { overloadIDs := c.checkedAST.GetOverloadIDs(e.ID()) if len(overloadIDs) == 0 { - return CostEstimate{} + var argCostSum CostEstimate + for _, a := range argCosts { + argCostSum = argCostSum.Add(a) + } + return FixedCostEstimate(1).Add(sum).Add(argCostSum) } var targetType *AstNode if call.IsMemberFunction() { @@ -429,13 +415,10 @@ func (c *coster) costCall(e ast.Expr) CostEstimate { switch overload { case overloads.IndexList: if len(args) > 0 { - // note: assigning resultSize here could be redundant with the path-based lookup later - resultSize = c.computeEntrySize(args[0]).valSize() c.addPath(e, append(c.getPath(args[0]), "@items")) } case overloads.IndexMap: if len(args) > 0 { - resultSize = c.computeEntrySize(args[0]).valSize() c.addPath(e, append(c.getPath(args[0]), "@values")) } } @@ -464,16 +447,9 @@ func (c *coster) maybeUnwrapDynCall(e ast.Expr) *CostEstimate { func (c *coster) costCreateList(e ast.Expr) CostEstimate { create := e.AsList() var sum CostEstimate - itemSize := SizeEstimate{Min: math.MaxUint64, Max: 0} - if create.Size() == 0 { - itemSize.Min = 0 - } - for _, e := range create.Elements() { - sum = sum.Add(c.cost(e)) - is := c.sizeOrUnknown(e) - itemSize = itemSize.Union(is) + for _, elem := range create.Elements() { + sum = sum.Add(c.cost(elem)) } - c.setEntrySize(e, &entrySizeEstimate{containerKind: types.ListKind, key: FixedSizeEstimate(1), val: itemSize}) return sum.Add(createListBaseCost) } @@ -481,24 +457,11 @@ func (c *coster) costCreateList(e ast.Expr) CostEstimate { func (c *coster) costCreateMap(e ast.Expr) CostEstimate { mapVal := e.AsMap() var sum CostEstimate - keySize := SizeEstimate{Min: math.MaxUint64, Max: 0} - valSize := SizeEstimate{Min: math.MaxUint64, Max: 0} - if mapVal.Size() == 0 { - valSize.Min = 0 - keySize.Min = 0 - } for _, ent := range mapVal.Entries() { entry := ent.AsMapEntry() sum = sum.Add(c.cost(entry.Key())) sum = sum.Add(c.cost(entry.Value())) - // Compute the key size range - ks := c.sizeOrUnknown(entry.Key()) - keySize = keySize.Union(ks) - // Compute the value size range - vs := c.sizeOrUnknown(entry.Value()) - valSize = valSize.Union(vs) - } - c.setEntrySize(e, &entrySizeEstimate{containerKind: types.MapKind, key: keySize, val: valSize}) + } return sum.Add(createMapBaseCost) } @@ -552,13 +515,16 @@ func (c *coster) costComprehension(e ast.Expr) CostEstimate { case ast.LiteralKind: c.setSize(e, c.computeSize(comp.AccuInit())) case ast.ListKind, ast.MapKind: - c.setSize(e, &rangeCnt) - // For a step which produces a container value, it will have an entry size associated - // with its expression id. - if stepEntrySize := c.computeEntrySize(comp.LoopStep()); stepEntrySize != nil { - c.setEntrySize(e, stepEntrySize) - break + accuSize := rangeCnt + if stepSize := c.computeSize(comp.LoopStep()); stepSize != nil { + if stepSize.Elem != nil { + accuSize.Elem = stepSize.Elem + } + if stepSize.Key != nil { + accuSize.Key = stepSize.Key + } } + c.setSize(e, &accuSize) } return sum } @@ -590,148 +556,47 @@ func (c *coster) costBind(e ast.Expr) CostEstimate { return sum } -// functionCost calculates the estimated call cost and result size for an overload invocation. -func (c *coster) functionCost(e ast.Expr, function, overloadID string, target *AstNode, args []AstNode, argCosts []CostEstimate) CallEstimate { - argCostSum := func() CostEstimate { - var sum CostEstimate - for _, a := range argCosts { - sum = sum.Add(a) +func calculateArgCost(overloadID string, argCosts []CostEstimate) CostEstimate { + switch overloadID { + case overloads.LogicalOr, overloads.LogicalAnd: + if len(argCosts) == 2 { + return CostEstimate{Min: argCosts[0].Min, Max: argCosts[0].Add(argCosts[1]).Max} + } + case overloads.Conditional: + if len(argCosts) == 3 { + return argCosts[0].Add(argCosts[1].Union(argCosts[2])) } - return sum } + var sum CostEstimate + for _, a := range argCosts { + sum = sum.Add(a) + } + return sum +} + +func (c *coster) functionCost(e ast.Expr, function, overloadID string, target *AstNode, args []AstNode, argCosts []CostEstimate) CallEstimate { + argCost := calculateArgCost(overloadID, argCosts) if len(c.overloadEstimators) != 0 { if estimator, found := c.overloadEstimators[overloadID]; found { if est := estimator(c.estimator, target, args); est != nil { - callEst := *est - return CallEstimate{CostEstimate: callEst.Add(argCostSum()), ResultSize: est.ResultSize} + return CallEstimate{CostEstimate: est.Add(argCost), ResultSize: est.ResultSize} } } } - if est := c.estimator.EstimateCallCost(function, overloadID, target, args); est != nil { - callEst := *est - return CallEstimate{CostEstimate: callEst.Add(argCostSum()), ResultSize: est.ResultSize} - } - switch overloadID { - // O(n) functions - case overloads.ExtFormatString: - if target != nil { - // ResultSize not calculated because we can't bound the max size. - return CallEstimate{ - CostEstimate: c.sizeOrUnknown(*target).MultiplyByCostFactor(StringTraversalCostFactor).Add(argCostSum())} - } - case overloads.StringToBytes: - if len(args) == 1 { - sz := c.sizeOrUnknown(args[0]) - // ResultSize max is when each char converts to 4 bytes. - return CallEstimate{ - CostEstimate: sz.MultiplyByCostFactor(StringTraversalCostFactor).Add(argCostSum()), - ResultSize: &SizeEstimate{Min: sz.Min, Max: sz.Max * 4}} - } - case overloads.BytesToString: - if len(args) == 1 { - sz := c.sizeOrUnknown(args[0]) - // ResultSize min is when 4 bytes convert to 1 char. - return CallEstimate{ - CostEstimate: sz.MultiplyByCostFactor(StringTraversalCostFactor).Add(argCostSum()), - ResultSize: &SizeEstimate{Min: sz.Min / 4, Max: sz.Max}} - } - case overloads.ExtQuoteString: - if len(args) == 1 { - sz := c.sizeOrUnknown(args[0]) - // ResultSize max is when each char is escaped. 2 quote chars always added. - return CallEstimate{ - CostEstimate: sz.MultiplyByCostFactor(StringTraversalCostFactor).Add(argCostSum()), - ResultSize: &SizeEstimate{Min: sz.Min + 2, Max: sz.Max*2 + 2}} - } - case overloads.StartsWithString, overloads.EndsWithString: - if len(args) == 1 { - return CallEstimate{CostEstimate: c.sizeOrUnknown(args[0]).MultiplyByCostFactor(StringTraversalCostFactor).Add(argCostSum())} - } - case overloads.InList: - // If a list is composed entirely of constant values this is O(1), but we don't account for that here. - // We just assume all list containment checks are O(n). - if len(args) == 2 { - return CallEstimate{CostEstimate: c.sizeOrUnknown(args[1]).MultiplyByCostFactor(1).Add(argCostSum())} + if c.estimator != nil { + if est := c.estimator.EstimateCallCost(function, overloadID, target, args); est != nil { + return CallEstimate{CostEstimate: est.Add(argCost), ResultSize: est.ResultSize} } - // O(nm) functions - case overloads.Matches, overloads.MatchesString: - // https://swtch.com/~rsc/regexp/regexp1.html applies to RE2 implementation supported by CEL - var strNode, regexNode AstNode - if overloadID == overloads.MatchesString && target != nil && len(args) == 1 { - strNode = *target - regexNode = args[0] - } else if overloadID == overloads.Matches && target == nil && len(args) == 2 { - strNode = args[0] - regexNode = args[1] - } - if strNode != nil && regexNode != nil { - // Add one to string length for purposes of cost calculation to prevent product of string and regex to be 0 - // in case where string is empty but regex is still expensive. - strCost := c.sizeOrUnknown(strNode).Add(SizeEstimate{Min: 1, Max: 1}).MultiplyByCostFactor(StringTraversalCostFactor) - // We don't know how many expressions are in the regex, just the string length (a huge - // improvement here would be to somehow get a count the number of expressions in the regex or - // how many states are in the regex state machine and use that to measure regex cost). - // For now, we're making a guess that each expression in a regex is typically at least 4 chars - // in length. - regexCost := c.sizeOrUnknown(regexNode).MultiplyByCostFactor(RegexStringLengthCostFactor) - return CallEstimate{CostEstimate: strCost.Multiply(regexCost).Add(argCostSum())} - } - case overloads.ContainsString: - if target != nil && len(args) == 1 { - strCost := c.sizeOrUnknown(*target).MultiplyByCostFactor(StringTraversalCostFactor) - substrCost := c.sizeOrUnknown(args[0]).MultiplyByCostFactor(StringTraversalCostFactor) - return CallEstimate{CostEstimate: strCost.Multiply(substrCost).Add(argCostSum())} - } - case overloads.LogicalOr, overloads.LogicalAnd: - lhs := argCosts[0] - rhs := argCosts[1] - // min cost is min of LHS for short circuited && or || - argCost := CostEstimate{Min: lhs.Min, Max: lhs.Add(rhs).Max} - return CallEstimate{CostEstimate: argCost} - case overloads.Conditional: - size := c.sizeOrUnknown(args[1]).Union(c.sizeOrUnknown(args[2])) - resultEntrySize := c.computeEntrySize(args[1].Expr()).union(c.computeEntrySize(args[2].Expr())) - c.setEntrySize(e, resultEntrySize) - conditionalCost := argCosts[0] - ifTrueCost := argCosts[1] - ifFalseCost := argCosts[2] - argCost := conditionalCost.Add(ifTrueCost.Union(ifFalseCost)) - return CallEstimate{CostEstimate: argCost, ResultSize: &size} - case overloads.AddString, overloads.AddBytes, overloads.AddList: - if len(args) == 2 { - lhsSize := c.sizeOrUnknown(args[0]) - rhsSize := c.sizeOrUnknown(args[1]) - resultSize := lhsSize.Add(rhsSize) - rhsEntrySize := c.computeEntrySize(args[0].Expr()) - lhsEntrySize := c.computeEntrySize(args[1].Expr()) - resultEntrySize := rhsEntrySize.union(lhsEntrySize) - if resultEntrySize != nil { - c.setEntrySize(e, resultEntrySize) - } - switch overloadID { - case overloads.AddList: - // list concatenation is O(1), but we handle it here to track size - return CallEstimate{CostEstimate: FixedCostEstimate(1).Add(argCostSum()), ResultSize: &resultSize} - default: - return CallEstimate{CostEstimate: resultSize.MultiplyByCostFactor(StringTraversalCostFactor).Add(argCostSum()), ResultSize: &resultSize} + } + if c.sizingStrategy != nil { + if estimator, found := c.getSizingOverloadEstimators()[overloadID]; found { + if est := estimator(c.estimator, target, args); est != nil { + return CallEstimate{CostEstimate: est.Add(argCost), ResultSize: est.ResultSize} } } - case overloads.LessString, overloads.GreaterString, overloads.LessEqualsString, overloads.GreaterEqualsString, - overloads.LessBytes, overloads.GreaterBytes, overloads.LessEqualsBytes, overloads.GreaterEqualsBytes, - overloads.Equals, overloads.NotEquals: - lhsCost := c.sizeOrUnknown(args[0]) - rhsCost := c.sizeOrUnknown(args[1]) - min := uint64(0) - smallestMax := lhsCost.Max - if rhsCost.Max < smallestMax { - smallestMax = rhsCost.Max - } - if smallestMax > 0 { - min = 1 - } - // equality of 2 scalar values results in a cost of 1 - return CallEstimate{ - CostEstimate: CostEstimate{Min: min, Max: smallestMax}.MultiplyByCostFactor(StringTraversalCostFactor).Add(argCostSum()), + } else if estimator, found := stdOverloadEstimators[overloadID]; found { + if est := estimator(c.estimator, target, args); est != nil { + return CallEstimate{CostEstimate: est.Add(argCost), ResultSize: est.ResultSize} } } // O(1) functions @@ -739,7 +604,7 @@ func (c *coster) functionCost(e ast.Expr, function, overloadID string, target *A // Benchmarks suggest that most of the other operations take +/- 50% of a base cost unit // which on an Intel xeon 2.20GHz CPU is 50ns. - return CallEstimate{CostEstimate: FixedCostEstimate(1).Add(argCostSum())} + return CallEstimate{CostEstimate: FixedCostEstimate(1).Add(argCost)} } // getType returns the deduced type of an expression ID from the checked AST. @@ -749,12 +614,25 @@ func (c *coster) getType(e ast.Expr) *types.Type { // getPath returns the tracked field path for an expression node, resolving through local variables if needed. func (c *coster) getPath(e ast.Expr) []string { + if e == nil { + return nil + } if e.Kind() == ast.IdentKind { if v, found := c.peekLocalVar(e.AsIdent()); found { return v.path[:] } } - return c.exprPaths[e.ID()][:] + if path, ok := c.exprPaths[e.ID()]; ok { + return path + } + return nil +} + +func (c *coster) getSizingOverloadEstimators() map[string]FunctionEstimator { + if c.sizingOverloadEstimators == nil { + c.sizingOverloadEstimators = StandardOverloadEstimatorsWithOptions(c.sizingStrategy) + } + return c.sizingOverloadEstimators } // addPath associates an expression ID with its path. @@ -762,19 +640,17 @@ func (c *coster) addPath(e ast.Expr, path []string) { c.exprPaths[e.ID()] = path } +var ( + stdOverloadEstimators = StandardOverloadEstimators() +) + func isAccumulatorVar(name string) bool { return name == accumulatorName || name == hiddenAccumulatorName } -// newAstNode creates an AstNode from an expression with path, type, and computed size. -func (c *coster) newAstNode(e ast.Expr) *astNode { - path := c.getPath(e) - if len(path) > 0 && isAccumulatorVar(path[0]) { - // only provide paths to root vars; omit accumulator vars - path = nil - } - return &astNode{ - path: path, +func (c *coster) newAstNode(e ast.Expr) AstNode { + return astNode{ + path: c.getPath(e), t: c.getType(e), expr: e, derivedSize: c.computeSize(e)} @@ -789,25 +665,99 @@ func (c *coster) setSize(e ast.Expr, size *SizeEstimate) { c.computedSizes[e.ID()] = *size } -// sizeOrUnknown extracts the size estimate from an ast.Expr or AstNode, falling back to UnknownSizeEstimate. -func (c *coster) sizeOrUnknown(node any) SizeEstimate { - switch v := node.(type) { - case ast.Expr: - if sz := c.computeSize(v); sz != nil { +func (c *coster) sizeOrUnknown(node ast.Expr) SizeEstimate { + if sz := c.computeSize(node); sz != nil { + return *sz + } + return UnknownSizeEstimate() +} + +// copySizeEstimates copies computed sizes and entry sizes from a source expression to a destination expression. +func (c *coster) copySizeEstimates(dst, src ast.Expr) { + c.setSize(dst, c.computeSize(src)) +} + +func (c *coster) getSizingStrategy() SizingStrategy { + if c.sizingStrategy != nil { + return c.sizingStrategy + } + return defaultSizing +} + +type estimatorContext struct { + coster *coster + estimator Estimator + target *AstNode + args []AstNode +} + +func (e *estimatorContext) Estimator() Estimator { + return e.estimator +} + +func (e *estimatorContext) Arg(index int) (SizeEstimate, bool) { + if index < len(e.args) { + return e.Size(e.args[index]), true + } + return UnknownSizeEstimate(), false +} + +func (e *estimatorContext) Target() (SizeEstimate, bool) { + if e.target != nil { + return e.Size(*e.target), true + } + return UnknownSizeEstimate(), false +} + +func (e *estimatorContext) Result() (SizeEstimate, bool) { + return UnknownSizeEstimate(), false +} + +func (e *estimatorContext) TargetType() (*types.Type, bool) { + if e.target != nil && (*e.target) != nil { + return (*e.target).Type(), true + } + return nil, false +} + +func (e *estimatorContext) ArgType(index int) (*types.Type, bool) { + if index < len(e.args) && e.args[index] != nil { + return e.args[index].Type(), true + } + return nil, false +} + +func (e *estimatorContext) Size(node AstNode) SizeEstimate { + if node == nil { + return UnknownSizeEstimate() + } + if sz := node.ComputedSize(); sz != nil { + return *sz + } + if e.coster != nil && node.Expr() != nil { + if sz := e.coster.computeSize(node.Expr()); sz != nil { return *sz } - case AstNode: - if sz := v.ComputedSize(); sz != nil { + } + if e.coster != nil { + if sz, ok := e.coster.getSizingStrategy().EstimateSize(e, node); ok { + return sz + } + } else if e.estimator != nil { + if sz := e.estimator.EstimateSize(node); sz != nil { return *sz } } return UnknownSizeEstimate() } -// copySizeEstimates copies computed sizes and entry sizes from a source expression to a destination expression. -func (c *coster) copySizeEstimates(dst, src ast.Expr) { - c.setSize(dst, c.computeSize(src)) - c.setEntrySize(dst, c.computeEntrySize(src)) +func (c *coster) newEstimateContext(target *AstNode, args []AstNode) *estimatorContext { + return &estimatorContext{ + coster: c, + estimator: c.estimator, + target: target, + args: args, + } } // computeSize resolves the size estimate for an expression, caching the result when found. @@ -818,44 +768,20 @@ func (c *coster) computeSize(e ast.Expr) *SizeEstimate { if size := computeExprSize(e); size != nil { return size } - // Ensure size estimates are computed first as users may choose to override the costs that - // CEL would otherwise ascribe to the type. - node := astNode{expr: e, path: c.getPath(e), t: c.getType(e)} - if size := c.estimator.EstimateSize(node); size != nil { - // storing the computed size should reduce calls to EstimateSize() - c.computedSizes[e.ID()] = *size - return size - } - if size := computeTypeSize(c.getType(e)); size != nil { - return size - } if e.Kind() == ast.IdentKind { varName := e.AsIdent() if v, ok := c.peekLocalVar(varName); ok && v.size != nil { return v.size } } - return nil -} - -// setEntrySize associates an expression with its container entry size estimate. -func (c *coster) setEntrySize(e ast.Expr, size *entrySizeEstimate) { - if size == nil { - return - } - c.computedEntrySizes[e.ID()] = *size -} - -// computeEntrySize looks up or resolves the container entry size estimate for an expression. -func (c *coster) computeEntrySize(e ast.Expr) *entrySizeEstimate { - if sz, found := c.computedEntrySizes[e.ID()]; found { - return &sz + node := astNode{expr: e, path: c.getPath(e), t: c.getType(e)} + ctx := c.newEstimateContext(nil, nil) + if size, ok := c.getSizingStrategy().EstimateSize(ctx, node); ok { + c.computedSizes[e.ID()] = size + return &size } - if e.Kind() == ast.IdentKind { - varName := e.AsIdent() - if v, ok := c.peekLocalVar(varName); ok && v.entrySize != nil { - return v.entrySize - } + if size := computeTypeSize(c.getType(e)); size != nil { + return size } return nil } @@ -880,10 +806,6 @@ func computeExprSize(expr ast.Expr) *SizeEstimate { default: return nil } - case ast.ListKind: - v = uint64(expr.AsList().Size()) - case ast.MapKind: - v = uint64(expr.AsMap().Size()) default: return nil } @@ -905,7 +827,7 @@ func computeTypeSize(t *types.Type) *SizeEstimate { // in addition to protobuf.Any and protobuf.Value (their size is not knowable at compile time). func isScalar(t *types.Type) bool { switch t.Kind() { - case types.BoolKind, types.DoubleKind, types.DurationKind, types.IntKind, types.TimestampKind, types.UintKind: + case types.BoolKind, types.DoubleKind, types.DurationKind, types.IntKind, types.TimestampKind, types.UintKind, types.TypeKind: return true case types.OpaqueKind: if t.TypeName() == "optional_type" { diff --git a/common/cost/estimator_test.go b/common/cost/estimator_test.go index 5bde4c15..835ae5b3 100644 --- a/common/cost/estimator_test.go +++ b/common/cost/estimator_test.go @@ -15,18 +15,23 @@ package cost_test import ( + "fmt" "math" "strings" "testing" "cel.dev/cel-go/checker" "cel.dev/cel-go/common" + "cel.dev/cel-go/common/ast" "cel.dev/cel-go/common/containers" "cel.dev/cel-go/common/cost" "cel.dev/cel-go/common/decls" + "cel.dev/cel-go/common/operators" "cel.dev/cel-go/common/overloads" "cel.dev/cel-go/common/stdlib" "cel.dev/cel-go/common/types" + "cel.dev/cel-go/common/types/ref" + "cel.dev/cel-go/common/types/traits" "cel.dev/cel-go/parser" proto3pb "cel.dev/cel-go/test/proto3pb" @@ -48,7 +53,7 @@ func TestCost(t *testing.T) { expr string vars []*decls.VariableDecl hints map[string]uint64 - options []cost.CostOption + options []cost.Option wanted cost.CostEstimate }{ { @@ -80,7 +85,7 @@ func TestCost(t *testing.T) { expr: `has(input.single_int32)`, vars: []*decls.VariableDecl{decls.NewVariable("input", types.NewObjectType("google.expr.proto3.test.TestAllTypes"))}, wanted: cost.CostEstimate{Min: 1, Max: 1}, - options: []cost.CostOption{cost.PresenceTestHasCost(false)}, + options: []cost.Option{cost.PresenceTestHasCost(false)}, }, { name: "select: field test only", @@ -93,14 +98,14 @@ func TestCost(t *testing.T) { expr: `has(input.testAttr.nestedAttr)`, vars: []*decls.VariableDecl{decls.NewVariable("input", nestedMap)}, wanted: cost.CostEstimate{Min: 3, Max: 3}, - options: []cost.CostOption{cost.PresenceTestHasCost(true)}, + options: []cost.Option{cost.PresenceTestHasCost(true)}, }, { name: "select: non-proto field test no has() cost", expr: `has(input.testAttr.nestedAttr)`, vars: []*decls.VariableDecl{decls.NewVariable("input", nestedMap)}, wanted: cost.CostEstimate{Min: 2, Max: 2}, - options: []cost.CostOption{cost.PresenceTestHasCost(false)}, + options: []cost.Option{cost.PresenceTestHasCost(false)}, }, { name: "select: non-proto field test", @@ -445,7 +450,7 @@ func TestCost(t *testing.T) { decls.NewVariable("str2", types.StringType), }, hints: map[string]uint64{"str1": 10, "str2": 10}, - options: []cost.CostOption{ + options: []cost.Option{ cost.OverloadCostEstimate(overloads.ContainsString, func(estimator cost.Estimator, target *cost.AstNode, args []cost.AstNode) *cost.CallEstimate { if target != nil && len(args) == 1 { @@ -577,17 +582,17 @@ func TestCost(t *testing.T) { { name: ".map list literal selection", expr: `[1,2,3,4,5].map(x, x)[4]`, - wanted: cost.CostEstimate{Min: 87, Max: 87}, + wanted: cost.CostEstimate{Min: 88, Max: 88}, }, { name: "nested array selection", expr: `[[1,2],[1,2],[1,2],[1,2],[1,2]][4]`, - wanted: cost.CostEstimate{Min: 61, Max: 61}, + wanted: cost.CostEstimate{Min: 62, Max: 62}, }, { name: "nested map selection", expr: `{'a': [1,2], 'b': [1,2], 'c': [1,2], 'd': [1,2], 'e': [1,2]}.b`, - wanted: cost.CostEstimate{Min: 81, Max: 81}, + wanted: cost.CostEstimate{Min: 82, Max: 82}, }, { name: "comprehension on nested list", @@ -632,12 +637,12 @@ func TestCost(t *testing.T) { { name: "literal map access", expr: `{'hello': 'hi'}['hello'] != {'hello': 'bye'}['hello']`, - wanted: cost.CostEstimate{Min: 63, Max: 63}, + wanted: cost.CostEstimate{Min: 65, Max: 65}, }, { name: "literal list access", expr: `['hello', 'hi'][0] != ['hello', 'bye'][1]`, - wanted: cost.CostEstimate{Min: 23, Max: 23}, + wanted: cost.CostEstimate{Min: 25, Max: 25}, }, { name: "type call", @@ -658,17 +663,17 @@ func TestCost(t *testing.T) { vars: []*decls.VariableDecl{ decls.NewVariable("self", types.NewMapType(types.StringType, types.IntType)), }, - wanted: cost.CostEstimate{Min: 5, Max: 1844674407370955268}, + wanted: cost.CostEstimate{Min: 5, Max: 5}, }, { name: "type literal equality cost", expr: `type(1) == int`, - wanted: cost.CostEstimate{Min: 3, Max: 1844674407370955266}, + wanted: cost.CostEstimate{Min: 3, Max: 3}, }, { name: "type variable equality cost", expr: `type(1) == int`, - wanted: cost.CostEstimate{Min: 3, Max: 1844674407370955266}, + wanted: cost.CostEstimate{Min: 3, Max: 3}, }, { name: "namespace variable equality", @@ -729,7 +734,7 @@ func TestCost(t *testing.T) { { name: "bytes list max", expr: "[bytes('012345678901'), bytes('012345678901'), bytes('012345678901'), bytes('012345678901'), bytes('012345678901')].max()", - options: []cost.CostOption{ + options: []cost.Option{ cost.OverloadCostEstimate("list_bytes_max", func(estimator cost.Estimator, target *cost.AstNode, args []cost.AstNode) *cost.CallEstimate { if target != nil { @@ -751,6 +756,154 @@ func TestCost(t *testing.T) { }, wanted: cost.CostEstimate{Min: 25, Max: 35}, }, + // cel.bind test cases + { + name: "bind: literal init and scalar result", + expr: `cel.bind(a, 'hello', a + '!')`, + wanted: cost.CostEstimate{Min: 12, Max: 12}, + }, + { + name: "bind: nested binds", + expr: `cel.bind(a, 'hello!', cel.bind(b, 'goodbye', a + ' and, ' + b))`, + wanted: cost.CostEstimate{Min: 26, Max: 26}, + }, + { + name: "bind: shadowed bind", + expr: `cel.bind(a, cel.bind(a, 'world', a + '!'), 'hello ' + a)`, + wanted: cost.CostEstimate{Min: 25, Max: 25}, + }, + { + name: "bind: with variable list and index", + expr: `cel.bind(a, input, a[0])`, + vars: []*decls.VariableDecl{decls.NewVariable("input", intList)}, + wanted: cost.CostEstimate{Min: 13, Max: 13}, + }, + { + name: "bind: with variable map and index", + expr: `cel.bind(m, input, m['key'])`, + vars: []*decls.VariableDecl{decls.NewVariable("input", types.NewMapType(types.StringType, types.StringType))}, + wanted: cost.CostEstimate{Min: 13, Max: 13}, + }, + { + name: "bind: with comprehension and size hints", + vars: []*decls.VariableDecl{decls.NewVariable("input", allList)}, + hints: map[string]uint64{"input": 100}, + expr: `cel.bind(a, input, a.all(x, true))`, + wanted: cost.CostEstimate{Min: 13, Max: 313}, + }, + { + name: "bind: nested with list and size hints", + vars: []*decls.VariableDecl{decls.NewVariable("input", nestedList)}, + hints: map[string]uint64{"input": 50, "input.@items": 10}, + expr: `cel.bind(a, input, a.all(x, x.all(y, true)))`, + wanted: cost.CostEstimate{Min: 13, Max: 1763}, + }, + { + name: "bind: unused bind variable", + expr: `cel.bind(a, [1, 2, 3], 42)`, + wanted: cost.CostEstimate{Min: 20, Max: 20}, + }, + { + name: "bind: derived size propagation to comprehension", + expr: `cel.bind(v, [1, 2, 3], v.all(x, true))`, + wanted: cost.CostEstimate{Min: 31, Max: 31}, + }, + + // Two-variable comprehension test cases + { + name: "two-var all: list literal", + expr: `[1, 2, 3].all(i, v, i < v)`, + wanted: cost.CostEstimate{Min: 20, Max: 29}, + }, + { + name: "two-var all: list variable with hints", + vars: []*decls.VariableDecl{decls.NewVariable("input", allList)}, + hints: map[string]uint64{"input": 100}, + expr: `input.all(i, v, true)`, + wanted: cost.CostEstimate{Min: 2, Max: 302}, + }, + { + name: "two-var all: map literal", + expr: `{"a": 1, "b": 2}.all(k, v, k != "" && v > 0)`, + wanted: cost.CostEstimate{Min: 37, Max: 43}, + }, + { + name: "two-var all: map variable with hints", + vars: []*decls.VariableDecl{decls.NewVariable("input", allMap)}, + hints: map[string]uint64{"input": 50}, + expr: `input.all(k, v, true)`, + wanted: cost.CostEstimate{Min: 2, Max: 152}, + }, + { + name: "two-var exists: list literal", + expr: `[1, 2, 3].exists(i, v, i == 1 && v == 2)`, + wanted: cost.CostEstimate{Min: 23, Max: 35}, + }, + { + name: "two-var exists: map literal", + expr: `{"a": 1, "b": 2}.exists(k, v, k == "a" && v == 1)`, + wanted: cost.CostEstimate{Min: 39, Max: 47}, + }, + { + name: "two-var existsOne: list literal", + expr: `[1, 2, 3].existsOne(i, v, v == 1)`, + wanted: cost.CostEstimate{Min: 21, Max: 24}, + }, + { + name: "two-var exists_one: list literal", + expr: `[1, 2, 3].exists_one(i, v, v == 1)`, + wanted: cost.CostEstimate{Min: 21, Max: 24}, + }, + { + name: "two-var transformList: 3-arg list literal", + expr: `[1, 2, 3].transformList(i, v, i + v)`, + wanted: cost.CostEstimate{Min: 66, Max: 66}, + }, + { + name: "two-var transformList: 4-arg with filter list literal", + expr: `[1, 2, 3].transformList(i, v, i % 2 == 0, i + v)`, + wanted: cost.CostEstimate{Min: 33, Max: 75}, + }, + { + name: "two-var transformList: 3-arg map literal", + expr: `{"a": 1, "b": 2}.transformList(k, v, k)`, + wanted: cost.CostEstimate{Min: 67, Max: 67}, + }, + { + name: "two-var transformMap: 3-arg map literal", + expr: `{"a": 1, "b": 2}.transformMap(k, v, v + 1)`, + wanted: cost.CostEstimate{Min: 71, Max: 71}, + }, + { + name: "two-var transformMap: 4-arg with filter map literal", + expr: `{"a": 1, "b": 2}.transformMap(k, v, v > 1, v + 1)`, + wanted: cost.CostEstimate{Min: 67, Max: 75}, + }, + { + name: "two-var transformMapEntry: 3-arg map literal", + expr: `{"a": 1, "b": 2}.transformMapEntry(k, v, {v: k})`, + wanted: cost.CostEstimate{Min: 129, Max: 129}, + }, + { + name: "two-var transformMapEntry: 4-arg with filter map literal", + expr: `{"a": 1, "b": 2}.transformMapEntry(k, v, v > 1, {v: k})`, + wanted: cost.CostEstimate{Min: 67, Max: 133}, + }, + { + name: "two-var nested all", + expr: `[1, 2].all(i, v, [1, 2].all(j, w, i + j < v + w))`, + wanted: cost.CostEstimate{Min: 17, Max: 79}, + }, + { + name: "bind with two-var comprehension", + expr: `cel.bind(l, [1, 2, 3], l.all(i, v, i < v))`, + wanted: cost.CostEstimate{Min: 31, Max: 40}, + }, + { + name: "bind with two-var transformList", + expr: `cel.bind(m, {"a": 1, "b": 2}, m.transformList(k, v, k))`, + wanted: cost.CostEstimate{Min: 78, Max: 78}, + }, } for _, tst := range cases { @@ -759,7 +912,7 @@ func TestCost(t *testing.T) { if tc.hints == nil { tc.hints = map[string]uint64{} } - p, err := parser.NewParser(parser.Macros(parser.AllMacros...)) + p, err := parser.NewParser(parser.Macros(testMacros...)) if err != nil { t.Fatalf("parser.NewParser() failed: %v", err) } @@ -785,7 +938,7 @@ func TestCost(t *testing.T) { decls.MemberOverload("list_bytes_max", []*types.Type{types.NewListType(types.BytesType)}, types.BytesType)) - err = e.AddFunctions(maxFunc) + err = e.AddFunctions(maxFunc, mapInsertFunctionDecl()) if err != nil { t.Fatalf("environment creation error: %v", err) } @@ -869,3 +1022,337 @@ func sizeEstimate(estimator cost.Estimator, t cost.AstNode) cost.SizeEstimate { } return cost.SizeEstimate{Min: 0, Max: math.MaxUint64} } + +type testCustomSizingStrategy struct{} + +func (testCustomSizingStrategy) EstimateSize(ctx cost.EstimateContext, node cost.AstNode) (cost.SizeEstimate, bool) { + if node.Path() != nil && len(node.Path()) > 0 && node.Path()[0] == "custom_str" { + return cost.SizeEstimate{Min: 10, Max: 20}, true + } + if node.Path() != nil && len(node.Path()) > 0 && node.Path()[0] == "custom_list" { + return cost.SizeEstimate{Min: 1, Max: 5, Elem: &cost.SizeEstimate{Min: 15, Max: 30}}, true + } + return cost.SizeEstimate{}, false +} + +func (testCustomSizingStrategy) TrackSize(ctx cost.TrackContext, value ref.Val) (uint64, bool) { + return cost.ActualSize(value), true +} + +func TestCustomSizingStrategy(t *testing.T) { + prse, err := parser.NewParser(parser.Macros(parser.AllMacros...)) + if err != nil { + t.Fatalf("parser.NewParser() failed: %v", err) + } + src := common.NewStringSource("custom_str.contains('abc')", "") + pe, errs := prse.Parse(src) + if len(errs.GetErrors()) != 0 { + t.Fatalf("Parse() failed: %v", errs.ToDisplayString()) + } + reg, err := types.NewRegistry() + if err != nil { + t.Fatalf("types.NewRegistry() failed: %v", err) + } + e, err := checker.NewEnv(containers.DefaultContainer, reg) + if err != nil { + t.Fatalf("checker.NewEnv() failed: %v", err) + } + err = e.AddFunctions(stdlib.Functions()...) + if err != nil { + t.Fatalf("AddFunctions failed: %v", err) + } + err = e.AddIdents(decls.NewVariable("custom_str", types.StringType)) + if err != nil { + t.Fatalf("AddIdents failed: %v", err) + } + checked, errs := checker.Check(pe, src, e) + if len(errs.GetErrors()) != 0 { + t.Fatalf("Check() failed: %v", errs.ToDisplayString()) + } + + res, err := cost.Cost(checked, nil, cost.EstimateSizingStrategy(testCustomSizingStrategy{})) + if err != nil { + t.Fatalf("Cost() failed: %v", err) + } + // 'abc' has length 3, cost traversal factor 0.1 -> ceil(3 * 0.1) = 1 + // custom_str has min 10, max 20 -> min ceil(10 * 0.1) = 1, max ceil(20 * 0.1) = 2 + // contains cost: min 1 * 1 = 1, max 2 * 1 = 2 + // ident cost = 1 + // total = ident(1) + call(min 1, max 2) = min 2, max 3 + if res.Min != 2 || res.Max != 3 { + t.Errorf("got cost %v, wanted {Min: 2, Max: 3}", res) + } +} + +var ( + testMacros = append( + append([]parser.Macro{}, parser.AllMacros...), + testCelBindMacro(), + testTwoVarAllMacro(), + testTwoVarExistsMacro(), + testTwoVarExistsOneMacro(), + testTwoVarExistsOneMacroNew(), + testTwoVarTransformListMacro(), + testTwoVarTransformListFilterMacro(), + testTwoVarTransformMapMacro(), + testTwoVarTransformMapFilterMacro(), + testTwoVarTransformMapEntryMacro(), + testTwoVarTransformMapEntryFilterMacro(), + ) +) + +func testCelBindMacro() parser.Macro { + return parser.NewReceiverMacro("bind", 3, func(eh parser.ExprHelper, target ast.Expr, args []ast.Expr) (ast.Expr, *common.Error) { + if target == nil || target.Kind() != ast.IdentKind || target.AsIdent() != "cel" { + return nil, nil + } + varIdent := args[0] + if varIdent.Kind() != ast.IdentKind { + return nil, eh.NewError(varIdent.ID(), "cel.bind() variable names must be simple identifiers") + } + varName := varIdent.AsIdent() + varInit := args[1] + resultExpr := args[2] + return eh.NewComprehension( + eh.NewList(), + "#unused", + varName, + varInit, + eh.NewLiteral(types.False), + eh.NewIdent(varName), + resultExpr, + ), nil + }) +} + +func testTwoVarAllMacro() parser.Macro { + return parser.NewReceiverMacro("all", 3, func(eh parser.ExprHelper, target ast.Expr, args []ast.Expr) (ast.Expr, *common.Error) { + iterVar1, iterVar2, err := extractTwoVarIterVars(eh, args[0], args[1]) + if err != nil { + return nil, err + } + return eh.NewComprehensionTwoVar( + target, + iterVar1, + iterVar2, + eh.AccuIdentName(), + eh.NewLiteral(types.True), + eh.NewCall(operators.NotStrictlyFalse, eh.NewAccuIdent()), + eh.NewCall(operators.LogicalAnd, eh.NewAccuIdent(), args[2]), + eh.NewAccuIdent(), + ), nil + }) +} + +func testTwoVarExistsMacro() parser.Macro { + return parser.NewReceiverMacro("exists", 3, func(eh parser.ExprHelper, target ast.Expr, args []ast.Expr) (ast.Expr, *common.Error) { + iterVar1, iterVar2, err := extractTwoVarIterVars(eh, args[0], args[1]) + if err != nil { + return nil, err + } + return eh.NewComprehensionTwoVar( + target, + iterVar1, + iterVar2, + eh.AccuIdentName(), + eh.NewLiteral(types.False), + eh.NewCall(operators.NotStrictlyFalse, eh.NewCall(operators.LogicalNot, eh.NewAccuIdent())), + eh.NewCall(operators.LogicalOr, eh.NewAccuIdent(), args[2]), + eh.NewAccuIdent(), + ), nil + }) +} + +func testTwoVarExistsOneMacro() parser.Macro { + return parser.NewReceiverMacro("exists_one", 3, testTwoVarExistsOneExpander) +} + +func testTwoVarExistsOneMacroNew() parser.Macro { + return parser.NewReceiverMacro("existsOne", 3, testTwoVarExistsOneExpander) +} + +func testTwoVarExistsOneExpander(eh parser.ExprHelper, target ast.Expr, args []ast.Expr) (ast.Expr, *common.Error) { + iterVar1, iterVar2, err := extractTwoVarIterVars(eh, args[0], args[1]) + if err != nil { + return nil, err + } + return eh.NewComprehensionTwoVar( + target, + iterVar1, + iterVar2, + eh.AccuIdentName(), + eh.NewLiteral(types.Int(0)), + eh.NewLiteral(types.True), + eh.NewCall(operators.Conditional, args[2], + eh.NewCall(operators.Add, eh.NewAccuIdent(), eh.NewLiteral(types.Int(1))), + eh.NewAccuIdent()), + eh.NewCall(operators.Equals, eh.NewAccuIdent(), eh.NewLiteral(types.Int(1))), + ), nil +} + +func testTwoVarTransformListMacro() parser.Macro { + return parser.NewReceiverMacro("transformList", 3, testTwoVarTransformListExpander) +} + +func testTwoVarTransformListFilterMacro() parser.Macro { + return parser.NewReceiverMacro("transformList", 4, testTwoVarTransformListExpander) +} + +func testTwoVarTransformListExpander(eh parser.ExprHelper, target ast.Expr, args []ast.Expr) (ast.Expr, *common.Error) { + iterVar1, iterVar2, err := extractTwoVarIterVars(eh, args[0], args[1]) + if err != nil { + return nil, err + } + var transform, filter ast.Expr + if len(args) == 4 { + filter = args[2] + transform = args[3] + } else { + transform = args[2] + } + step := eh.NewCall(operators.Add, eh.NewAccuIdent(), eh.NewList(transform)) + if filter != nil { + step = eh.NewCall(operators.Conditional, filter, step, eh.NewAccuIdent()) + } + return eh.NewComprehensionTwoVar( + target, + iterVar1, + iterVar2, + eh.AccuIdentName(), + eh.NewList(), + eh.NewLiteral(types.True), + step, + eh.NewAccuIdent(), + ), nil +} + +func testTwoVarTransformMapMacro() parser.Macro { + return parser.NewReceiverMacro("transformMap", 3, testTwoVarTransformMapExpander) +} + +func testTwoVarTransformMapFilterMacro() parser.Macro { + return parser.NewReceiverMacro("transformMap", 4, testTwoVarTransformMapExpander) +} + +func testTwoVarTransformMapExpander(eh parser.ExprHelper, target ast.Expr, args []ast.Expr) (ast.Expr, *common.Error) { + iterVar1, iterVar2, err := extractTwoVarIterVars(eh, args[0], args[1]) + if err != nil { + return nil, err + } + var transform, filter ast.Expr + if len(args) == 4 { + filter = args[2] + transform = args[3] + } else { + transform = args[2] + } + step := eh.NewCall("cel.@mapInsert", eh.NewAccuIdent(), eh.NewIdent(iterVar1), transform) + if filter != nil { + step = eh.NewCall(operators.Conditional, filter, step, eh.NewAccuIdent()) + } + return eh.NewComprehensionTwoVar( + target, + iterVar1, + iterVar2, + eh.AccuIdentName(), + eh.NewMap(), + eh.NewLiteral(types.True), + step, + eh.NewAccuIdent(), + ), nil +} + +func testTwoVarTransformMapEntryMacro() parser.Macro { + return parser.NewReceiverMacro("transformMapEntry", 3, testTwoVarTransformMapEntryExpander) +} + +func testTwoVarTransformMapEntryFilterMacro() parser.Macro { + return parser.NewReceiverMacro("transformMapEntry", 4, testTwoVarTransformMapEntryExpander) +} + +func testTwoVarTransformMapEntryExpander(eh parser.ExprHelper, target ast.Expr, args []ast.Expr) (ast.Expr, *common.Error) { + iterVar1, iterVar2, err := extractTwoVarIterVars(eh, args[0], args[1]) + if err != nil { + return nil, err + } + var transform, filter ast.Expr + if len(args) == 4 { + filter = args[2] + transform = args[3] + } else { + transform = args[2] + } + step := eh.NewCall("cel.@mapInsert", eh.NewAccuIdent(), transform) + if filter != nil { + step = eh.NewCall(operators.Conditional, filter, step, eh.NewAccuIdent()) + } + return eh.NewComprehensionTwoVar( + target, + iterVar1, + iterVar2, + eh.AccuIdentName(), + eh.NewMap(), + eh.NewLiteral(types.True), + step, + eh.NewAccuIdent(), + ), nil +} + +func extractTwoVarIterVars(eh parser.ExprHelper, arg0, arg1 ast.Expr) (string, string, *common.Error) { + if arg0.Kind() != ast.IdentKind { + return "", "", eh.NewError(arg0.ID(), "argument must be a simple name") + } + if arg1.Kind() != ast.IdentKind { + return "", "", eh.NewError(arg1.ID(), "argument must be a simple name") + } + iterVar1 := arg0.AsIdent() + iterVar2 := arg1.AsIdent() + if iterVar1 == iterVar2 { + return "", "", eh.NewError(arg1.ID(), fmt.Sprintf("duplicate variable name: %s", iterVar1)) + } + if iterVar1 == eh.AccuIdentName() || iterVar1 == parser.AccumulatorName { + return "", "", eh.NewError(arg0.ID(), "iteration variable overwrites accumulator variable") + } + if iterVar2 == eh.AccuIdentName() || iterVar2 == parser.AccumulatorName { + return "", "", eh.NewError(arg1.ID(), "iteration variable overwrites accumulator variable") + } + return iterVar1, iterVar2, nil +} + +func mapInsertFunctionDecl() *decls.FunctionDecl { + kType := types.NewTypeParamType("K") + vType := types.NewTypeParamType("V") + mapKVType := types.NewMapType(kType, vType) + fn, _ := decls.NewFunction("cel.@mapInsert", + decls.Overload("@mapInsert_map_key_value", + []*types.Type{mapKVType, kType, vType}, + mapKVType), + decls.Overload("@mapInsert_map_map", + []*types.Type{mapKVType, mapKVType}, + mapKVType), + decls.SingletonFunctionBinding(func(args ...ref.Val) ref.Val { + if len(args) == 3 { + m := args[0].(traits.Mapper) + k := args[1] + v := args[2] + return types.InsertMapKeyValue(m, k, v) + } + if len(args) == 2 { + tm := args[0].(traits.Mapper) + um := args[1].(traits.Mapper) + umIt := um.Iterator() + for umIt.HasNext() == types.True { + k := umIt.Next() + updateOrErr := types.InsertMapKeyValue(tm, k, um.Get(k)) + if types.IsError(updateOrErr) { + return updateOrErr + } + tm = updateOrErr.(traits.Mapper) + } + return tm + } + return types.NoSuchOverloadErr() + }), + ) + return fn +} diff --git a/common/cost/tracker.go b/common/cost/tracker.go index 00e4f9b1..c576f200 100644 --- a/common/cost/tracker.go +++ b/common/cost/tracker.go @@ -15,7 +15,6 @@ package cost import ( - "cel.dev/cel-go/common/overloads" "cel.dev/cel-go/common/types/ref" ) @@ -82,6 +81,14 @@ func OverloadTracker(overloadID string, fnTracker FunctionTracker) TrackerOption } } +// TrackerSizingStrategy configures a SizingStrategy for runtime cost tracking. +func TrackerSizingStrategy(strategy SizingStrategy) TrackerOption { + return func(tracker *Tracker) error { + tracker.sizingStrategy = strategy + return nil + } +} + // LimitExceededError indicates that the actual cost limit was exceeded during evaluation. type LimitExceededError struct { Message string @@ -94,11 +101,13 @@ func (e LimitExceededError) Error() string { // Tracker represents the information needed for tracking runtime cost. type Tracker struct { - Estimator ActualCostEstimator - overloadTrackers map[string]FunctionTracker - Limit *uint64 - presenceTestHasCost bool - limitExceededHandler func() + Estimator ActualCostEstimator + overloadTrackers map[string]FunctionTracker + sizingStrategy SizingStrategy + sizingOverloadTrackers map[string]FunctionTracker + Limit *uint64 + presenceTestHasCost bool + limitExceededHandler func() cost uint64 } @@ -123,11 +132,13 @@ func NewTracker(estimator ActualCostEstimator, opts ...TrackerOption) (*Tracker, // The different clones can be used independently from each other. func (c *Tracker) Clone() (*Tracker, error) { tracker := &Tracker{ - Estimator: c.Estimator, - overloadTrackers: c.overloadTrackers, - Limit: c.Limit, - presenceTestHasCost: c.presenceTestHasCost, - limitExceededHandler: c.limitExceededHandler, + Estimator: c.Estimator, + overloadTrackers: c.overloadTrackers, + sizingStrategy: c.sizingStrategy, + sizingOverloadTrackers: c.sizingOverloadTrackers, + Limit: c.Limit, + presenceTestHasCost: c.presenceTestHasCost, + limitExceededHandler: c.limitExceededHandler, } return tracker, nil } @@ -213,6 +224,13 @@ func (c *Tracker) checkLimit() { } } +func (c *Tracker) getSizingOverloadTrackers() map[string]FunctionTracker { + if c.sizingOverloadTrackers == nil { + c.sizingOverloadTrackers = StandardOverloadTrackersWithOptions(c.sizingStrategy) + } + return c.sizingOverloadTrackers +} + // CostCall calculates the runtime cost for a function call. func (c *Tracker) CostCall(call Call, args []ref.Val, result ref.Val) uint64 { var total uint64 @@ -232,62 +250,31 @@ func (c *Tracker) CostCall(call Call, args []ref.Val, result ref.Val) uint64 { return total } } - // if user didn't specify, the default way of calculating runtime cost would be used. - // if user has their own implementation of ActualCostEstimator, make sure to cover the mapping between overloadId and cost calculation - switch call.OverloadID() { - // O(n) functions - case overloads.StartsWithString, overloads.EndsWithString: - total = SafeAdd(total, SafeMultiplyByFactor(ActualSize(args[1]), StringTraversalCostFactor)) - case overloads.StringToBytes, overloads.BytesToString, overloads.ExtQuoteString, overloads.ExtFormatString: - total = SafeAdd(total, SafeMultiplyByFactor(ActualSize(args[0]), StringTraversalCostFactor)) - case overloads.InList: - // If a list is composed entirely of constant values this is O(1), but we don't account for that here. - // We just assume all list containment checks are O(n). - total = SafeAdd(total, ActualSize(args[1])) - // O(min(m, n)) functions - case overloads.LessString, overloads.GreaterString, overloads.LessEqualsString, overloads.GreaterEqualsString, - overloads.LessBytes, overloads.GreaterBytes, overloads.LessEqualsBytes, overloads.GreaterEqualsBytes, - overloads.Equals, overloads.NotEquals: - // When we check the equality of 2 scalar values (e.g. 2 integers, 2 floating-point numbers, 2 booleans etc.), - // the CostTracker.ActualSize() function by definition returns 1 for each operand, resulting in an overall cost - // of 1. - lhsSize := ActualSize(args[0]) - rhsSize := ActualSize(args[1]) - minSize := min(rhsSize, lhsSize) - total = SafeAdd(total, SafeMultiplyByFactor(minSize, StringTraversalCostFactor)) - // O(m+n) functions - case overloads.AddString, overloads.AddBytes: - // In the worst case scenario, we would need to reallocate a new backing store and copy both operands over. - argSize := SafeAdd(ActualSize(args[0]), ActualSize(args[1])) - total = SafeAdd(total, SafeMultiplyByFactor(argSize, StringTraversalCostFactor)) - // O(nm) functions - case overloads.Matches, overloads.MatchesString: - // https://swtch.com/~rsc/regexp/regexp1.html applies to RE2 implementation supported by CEL - // Add one to string length for purposes of cost calculation to prevent product of string and regex to be 0 - // in case where string is empty but regex is still expensive. - strCost := SafeMultiplyByFactor(SafeAdd(1, ActualSize(args[0])), StringTraversalCostFactor) - // We don't know how many expressions are in the regex, just the string length (a huge - // improvement here would be to somehow get a count the number of expressions in the regex or - // how many states are in the regex state machine and use that to measure regex cost). - // For now, we're making a guess that each expression in a regex is typically at least 4 chars - // in length. - regexCost := SafeMultiplyByFactor(ActualSize(args[1]), RegexStringLengthCostFactor) - total = SafeAdd(total, SafeMultiply(strCost, regexCost)) - case overloads.ContainsString: - strCost := SafeMultiplyByFactor(ActualSize(args[0]), StringTraversalCostFactor) - substrCost := SafeMultiplyByFactor(ActualSize(args[1]), StringTraversalCostFactor) - total = SafeAdd(total, SafeMultiply(strCost, substrCost)) - - default: - // The following operations are assumed to have O(1) complexity. - // - AddList due to the implementation. Index lookup can be O(c) the - // number of concatenated lists, but we don't track that is cost calculations. - // - Conversions, since none perform a traversal of a type of unbound length. - // - Computing the size of strings, byte sequences, lists and maps. - // - Logical operations and all operators on fixed width scalars (comparisons, equality) - // - Any functions that don't have a declared cost either here or in provided ActualCostEstimator. - total = SafeAdd(total, 1) - + if c.sizingStrategy != nil { + if tracker, found := c.getSizingOverloadTrackers()[call.OverloadID()]; found { + callCost := tracker(args, result) + if callCost != nil { + total = SafeAdd(total, *callCost) + return total + } + } + } else if tracker, found := stdOverloadTrackers[call.OverloadID()]; found { + callCost := tracker(args, result) + if callCost != nil { + total = SafeAdd(total, *callCost) + return total + } } - return total + // The following operations are assumed to have O(1) complexity. + // - AddList due to the implementation. Index lookup can be O(c) the + // number of concatenated lists, but we don't track that in cost calculations. + // - Conversions, since none perform a traversal of a type of unbound length. + // - Computing the size of strings, byte sequences, lists and maps. + // - Logical operations and all operators on fixed width scalars (comparisons, equality) + // - Any functions that don't have a declared cost either here or in provided ActualCostEstimator. + return SafeAdd(total, 1) } + +var ( + stdOverloadTrackers = StandardOverloadTrackers() +) diff --git a/common/cost/tracker_test.go b/common/cost/tracker_test.go index 85f2b0e8..dac6fcc9 100644 --- a/common/cost/tracker_test.go +++ b/common/cost/tracker_test.go @@ -51,37 +51,63 @@ func (c testCall) OverloadID() string { return c.overloadID } -func TestCostTrackerBasic(t *testing.T) { +func TestCostTracker_BasicOperations(t *testing.T) { tracker, err := cost.NewTracker(nil, cost.TrackerPresenceTestHasCost(true), ) if err != nil { - t.Fatalf("NewCostTracker() failed: %v", err) + t.Fatalf("NewTracker() failed: %v", err) } - tracker.CreateList(1, nil) - if tracker.ActualCost() != cost.ListCreateBaseCost { - t.Errorf("ActualCost() after CreateList = %d, want %d", tracker.ActualCost(), cost.ListCreateBaseCost) - } - - tracker.CreateMap(2, nil) - if tracker.ActualCost() != cost.ListCreateBaseCost+cost.MapCreateBaseCost { - t.Errorf("ActualCost() after CreateMap = %d, want %d", tracker.ActualCost(), cost.ListCreateBaseCost+cost.MapCreateBaseCost) - } - - tracker.CreateStruct(3, nil) - if tracker.ActualCost() != cost.ListCreateBaseCost+cost.MapCreateBaseCost+cost.StructCreateBaseCost { - t.Errorf("ActualCost() after CreateStruct = %d, want %d", tracker.ActualCost(), cost.ListCreateBaseCost+cost.MapCreateBaseCost+cost.StructCreateBaseCost) - } - - tracker.EvalAttribute(4, false, nil) - if tracker.ActualCost() != cost.ListCreateBaseCost+cost.MapCreateBaseCost+cost.StructCreateBaseCost+cost.SelectAndIdentCost { - t.Errorf("ActualCost() after EvalAttribute = %d, want %d", tracker.ActualCost(), cost.ListCreateBaseCost+cost.MapCreateBaseCost+cost.StructCreateBaseCost+cost.SelectAndIdentCost) + tests := []struct { + name string + action func() + wantCost uint64 + }{ + { + name: "create_list", + action: func() { + tracker.CreateList(1, nil) + }, + wantCost: cost.ListCreateBaseCost, + }, + { + name: "create_map", + action: func() { + tracker.CreateMap(2, nil) + }, + wantCost: cost.ListCreateBaseCost + cost.MapCreateBaseCost, + }, + { + name: "create_struct", + action: func() { + tracker.CreateStruct(3, nil) + }, + wantCost: cost.ListCreateBaseCost + cost.MapCreateBaseCost + cost.StructCreateBaseCost, + }, + { + name: "eval_attribute", + action: func() { + tracker.EvalAttribute(4, false, nil) + }, + wantCost: cost.ListCreateBaseCost + cost.MapCreateBaseCost + cost.StructCreateBaseCost + cost.SelectAndIdentCost, + }, + { + name: "qualify", + action: func() { + tracker.Qualify(5) + }, + wantCost: cost.ListCreateBaseCost + cost.MapCreateBaseCost + cost.StructCreateBaseCost + cost.SelectAndIdentCost + 1, + }, } - tracker.Qualify(5) - if tracker.ActualCost() != cost.ListCreateBaseCost+cost.MapCreateBaseCost+cost.StructCreateBaseCost+cost.SelectAndIdentCost+1 { - t.Errorf("ActualCost() after Qualify = %d, want %d", tracker.ActualCost(), cost.ListCreateBaseCost+cost.MapCreateBaseCost+cost.StructCreateBaseCost+cost.SelectAndIdentCost+1) + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + tc.action() + if tracker.ActualCost() != tc.wantCost { + t.Errorf("ActualCost() = %d, want %d", tracker.ActualCost(), tc.wantCost) + } + }) } if !tracker.PresenceTestHasCost() { @@ -89,7 +115,7 @@ func TestCostTrackerBasic(t *testing.T) { } } -func TestCostTrackerLimit(t *testing.T) { +func TestCostTracker_LimitExceededPanic(t *testing.T) { var exceeded bool tracker, err := cost.NewTracker(nil, cost.TrackerLimit(15), @@ -98,7 +124,7 @@ func TestCostTrackerLimit(t *testing.T) { }), ) if err != nil { - t.Fatalf("NewCostTracker() failed: %v", err) + t.Fatalf("NewTracker() failed: %v", err) } tracker.CreateList(1, nil) // cost = 10 <= 15 @@ -119,7 +145,7 @@ func TestCostTrackerLimit(t *testing.T) { tracker.CreateList(2, nil) // cost = 20 > 15 -> panic } -func TestCostTrackerOverloadTracker(t *testing.T) { +func TestCostTracker_CustomOverloadTracker(t *testing.T) { tracker, err := cost.NewTracker(nil, cost.OverloadTracker("custom_op", func(args []ref.Val, result ref.Val) *uint64 { c := uint64(42) @@ -127,7 +153,7 @@ func TestCostTrackerOverloadTracker(t *testing.T) { }), ) if err != nil { - t.Fatalf("NewCostTracker() failed: %v", err) + t.Fatalf("NewTracker() failed: %v", err) } call := testCall{function: "custom", overloadID: "custom_op"} @@ -137,10 +163,10 @@ func TestCostTrackerOverloadTracker(t *testing.T) { } } -func TestCostTrackerClone(t *testing.T) { +func TestCostTracker_CloneStateIsolation(t *testing.T) { tracker, err := cost.NewTracker(nil) if err != nil { - t.Fatalf("NewCostTracker() failed: %v", err) + t.Fatalf("NewTracker() failed: %v", err) } tracker.Qualify(1) @@ -161,17 +187,68 @@ func TestCostTrackerClone(t *testing.T) { } } -func TestCostTrackerStandardFunctions(t *testing.T) { - tracker, err := cost.NewTracker(nil) - if err != nil { - t.Fatalf("NewCostTracker() failed: %v", err) +func TestCostTracker_StandardStringFunctionTracking(t *testing.T) { + adapter := types.DefaultTypeAdapter + + tests := []struct { + name string + overloadID string + function string + target ref.Val + arg ref.Val + result ref.Val + wantCost uint64 + }{ + { + name: "starts_with_string", + overloadID: overloads.StartsWithString, + function: "startsWith", + target: types.String("hello world"), + arg: types.String("hello"), // len 5 -> ceil(5 * 0.1) = 1 + result: types.True, + wantCost: 1, + }, + { + name: "ends_with_string", + overloadID: overloads.EndsWithString, + function: "endsWith", + target: types.String("hello world"), + arg: types.String("world"), // len 5 -> ceil(5 * 0.1) = 1 + result: types.True, + wantCost: 1, + }, + { + name: "contains_string", + overloadID: overloads.ContainsString, + function: "contains", + target: types.String("hello world"), + arg: types.String("lo wo"), // len 5 -> ceil(11*0.1) * ceil(5*0.1) = 2 * 1 = 2 + result: types.True, + wantCost: 2, + }, + { + name: "in_list_string", + overloadID: overloads.InList, + function: "@in", + target: types.String("item"), + arg: adapter.NativeToValue([]string{"a", "b", "c"}), + result: types.False, + wantCost: 3, + }, } - // StartsWith - tracker.EvalBinary(nil, 1, testCall{function: "startsWith", overloadID: overloads.StartsWithString}, types.String("hello world"), types.String("hello"), types.True) - // cost.ActualSize("hello") = 5. cost = ceil(5 * 0.1) = 1. - if tracker.ActualCost() != 1 { - t.Errorf("ActualCost() after startsWith = %d, want 1", tracker.ActualCost()) + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + tracker, err := cost.NewTracker(nil) + if err != nil { + t.Fatalf("NewTracker() failed: %v", err) + } + call := testCall{function: tc.function, overloadID: tc.overloadID} + tracker.EvalBinary(nil, 1, call, tc.target, tc.arg, tc.result) + if tracker.ActualCost() != tc.wantCost { + t.Errorf("ActualCost() = %d, want %d", tracker.ActualCost(), tc.wantCost) + } + }) } } @@ -253,7 +330,7 @@ func computeCost(t *testing.T, expr string, vars []*decls.VariableDecl, ctx inte t.Helper() s := common.NewTextSource(expr) - p, err := parser.NewParser(parser.Macros(parser.AllMacros...)) + p, err := parser.NewParser(parser.Macros(testMacros...)) if err != nil { t.Fatalf("Failed to initialize parser: %v", err) } @@ -266,6 +343,10 @@ func computeCost(t *testing.T, expr string, vars []*decls.VariableDecl, ctx inte reg := newTestRegistry(t, types.ProtoTypeDefs(&proto3pb.TestAllTypes{})) attrs := interpreter.NewAttributeFactory(cont, reg, reg) env := newTestEnv(t, cont, reg) + err = env.AddFunctions(mapInsertFunctionDecl()) + if err != nil { + t.Fatalf("Failed to add mapInsertFunctionDecl: %v", err) + } err = env.AddIdents(vars...) if err != nil { t.Fatalf("Failed to initialize env: %v", err) @@ -286,7 +367,7 @@ func computeCost(t *testing.T, expr string, vars []*decls.VariableDecl, ctx inte if err != nil { t.Fatalf("cost.Cost() failed: %v", err) } - interp := newStandardInterpreter(t, cont, reg, reg, attrs) + interp := newStandardInterpreter(t, cont, reg, reg, attrs, mapInsertFunctionDecl()) prg, err := interp.NewInterpretable(checked, interpreter.CostObserver(interpreter.CostTrackerFactory(func() (*cost.Tracker, error) { return costTracker, nil @@ -1007,6 +1088,167 @@ func TestRuntimeCost(t *testing.T) { expr: `[1,2,3].all(i, i in [1,2,3].map(j, j + j))`, want: 86, }, + // cel.bind runtime cost tracking test cases + { + name: "bind: literal init and scalar result", + expr: `cel.bind(a, 'hello', a + '!')`, + want: 12, + }, + { + name: "bind: nested binds", + expr: `cel.bind(a, 'hello!', cel.bind(b, 'goodbye', a + ' and, ' + b))`, + want: 26, + }, + { + name: "bind: shadowed bind", + expr: `cel.bind(a, cel.bind(a, 'world', a + '!'), 'hello ' + a)`, + want: 25, + }, + { + name: "bind: with variable list and index", + expr: `cel.bind(a, input, a[0])`, + vars: []*decls.VariableDecl{decls.NewVariable("input", intList)}, + want: 13, + in: map[string]any{"input": []int{1, 2}}, + }, + { + name: "bind: with variable map and index", + expr: `cel.bind(m, input, m['key'])`, + vars: []*decls.VariableDecl{decls.NewVariable("input", types.NewMapType(types.StringType, types.StringType))}, + want: 13, + in: map[string]any{"input": map[string]string{"key": "value"}}, + }, + { + name: "bind: with comprehension over empty list", + vars: []*decls.VariableDecl{decls.NewVariable("input", allList)}, + expr: `cel.bind(a, input, a.all(x, true))`, + want: 13, + in: map[string]any{ + "input": []*proto3pb.TestAllTypes{}, + }, + }, + { + name: "bind: with list and indexing", + expr: `cel.bind(a, [1, 2, 3], a[0])`, + want: 22, + }, + { + name: "bind: derived size propagation to comprehension", + expr: `cel.bind(v, [1, 2, 3], v.all(x, true))`, + want: 31, + }, + { + name: "bind: limit exceeded", + expr: `cel.bind(a, [1, 2, 3], a.all(x, true))`, + limit: 25, + expectExceedsLimit: true, + }, + + // Two-variable comprehension runtime cost tracking test cases + { + name: "two-var all: list literal", + expr: `[1, 2, 3].all(i, v, i < v)`, + want: 29, + }, + { + name: "two-var all: list variable early return false", + vars: []*decls.VariableDecl{decls.NewVariable("input", intList)}, + expr: `input.all(i, v, i > v) == false`, + want: 11, + in: map[string]any{"input": []int{1, 2, 3}}, + }, + { + name: "two-var all: list variable", + vars: []*decls.VariableDecl{decls.NewVariable("input", intList)}, + expr: `input.all(i, v, i < 5)`, + want: 17, + in: map[string]any{"input": []int{1, 2, 3}}, + }, + { + name: "two-var all: map literal early return false", + expr: `{'hello': 'world', 'taco': 'taco'}.all(k, v, k != v) == false`, + want: 44, + }, + { + name: "two-var exists: list literal", + expr: `[1, 2, 3].exists(i, v, i == 1 && v == 2)`, + want: 28, + }, + { + name: "two-var exists: map literal", + expr: `{"a": 1, "b": 2}.exists(k, v, k == "a" && v == 1)`, + want: 42, + }, + { + name: "two-var existsOne: list variable", + vars: []*decls.VariableDecl{decls.NewVariable("input", intList)}, + expr: `input.existsOne(i, v, v == 1)`, + want: 11, + in: map[string]any{"input": []int{1, 2, 3}}, + }, + { + name: "two-var exists_one: list variable", + vars: []*decls.VariableDecl{decls.NewVariable("input", intList)}, + expr: `input.exists_one(i, v, v == 1)`, + want: 11, + in: map[string]any{"input": []int{1, 2, 3}}, + }, + { + name: "two-var transformList: 3-arg list literal", + expr: `[1, 2, 3].transformList(i, v, i + v)`, + want: 66, + }, + { + name: "two-var transformList: 4-arg with filter list literal", + expr: `[1, 2, 3].transformList(i, v, i % 2 == 0, i + v)`, + want: 60, + }, + { + name: "two-var transformList: 3-arg map literal", + expr: `{"a": 1, "b": 2}.transformList(k, v, k)`, + want: 67, + }, + { + name: "two-var transformMap: 3-arg map literal", + expr: `{"a": 1, "b": 2}.transformMap(k, v, v + 1)`, + want: 71, + }, + { + name: "two-var transformMap: 4-arg with filter map literal", + expr: `{"a": 1, "b": 2}.transformMap(k, v, v > 1, v + 1)`, + want: 70, + }, + { + name: "two-var transformMapEntry: 3-arg map literal", + expr: `{"a": 1, "b": 2}.transformMapEntry(k, v, {v: k})`, + want: 129, + }, + { + name: "two-var transformMapEntry: 4-arg with filter map literal", + expr: `{"a": 1, "b": 2}.transformMapEntry(k, v, v > 1, {v: k})`, + want: 99, + }, + { + name: "two-var nested all", + expr: `[1, 2].all(i, v, [1, 2].all(j, w, i + j < v + w))`, + want: 79, + }, + { + name: "bind with two-var comprehension", + expr: `cel.bind(l, [1, 2, 3], l.all(i, v, i < v))`, + want: 40, + }, + { + name: "bind with two-var transformList", + expr: `cel.bind(m, {"a": 1, "b": 2}, m.transformList(k, v, k))`, + want: 78, + }, + { + name: "two-var transformList: limit exceeded", + expr: `[1, 2, 3, 4, 5].transformList(i, v, i + v)`, + limit: 50, + expectExceedsLimit: true, + }, } for _, tc := range cases { From 4143b690d6318646efe88d025d0820616b9ebe3d Mon Sep 17 00:00:00 2001 From: TristonianJones Date: Thu, 10 Sep 2026 11:21:04 -0700 Subject: [PATCH 3/6] Updates post-merge, simplifications to sizing strategy setup --- common/cost/BUILD.bazel | 2 -- common/cost/estimator.go | 20 ++++++++++++-------- common/cost/model.go | 10 ++++++++-- common/cost/standard.go | 10 ++++++++-- common/cost/tracker.go | 22 ++++++++++++---------- interpreter/interpretable.go | 4 ++-- 6 files changed, 42 insertions(+), 26 deletions(-) diff --git a/common/cost/BUILD.bazel b/common/cost/BUILD.bazel index 38cbbcf5..b8e85e0e 100644 --- a/common/cost/BUILD.bazel +++ b/common/cost/BUILD.bazel @@ -8,7 +8,6 @@ package( go_library( name = "go_default_library", srcs = [ - "aggregate_strategy.go", "cost.go", "default_strategy.go", "estimator.go", @@ -32,7 +31,6 @@ go_test( name = "go_default_test", size = "small", srcs = [ - "aggregate_strategy_test.go", "cost_test.go", "default_strategy_test.go", "estimator_test.go", diff --git a/common/cost/estimator.go b/common/cost/estimator.go index 0f4ad355..7260ef19 100644 --- a/common/cost/estimator.go +++ b/common/cost/estimator.go @@ -135,6 +135,9 @@ func OverloadCostEstimate(overloadID string, functionCoster FunctionEstimator) O // EstimateSizingStrategy configures a custom SizingStrategy for cost estimation. func EstimateSizingStrategy(strategy SizingStrategy) Option { return func(c *coster) error { + if strategy == nil { + strategy = DefaultSizingStrategy() + } c.sizingStrategy = strategy return nil } @@ -146,6 +149,7 @@ func Cost(checked *ast.AST, estimator Estimator, opts ...Option) (CostEstimate, checkedAST: checked, estimator: estimator, overloadEstimators: map[string]FunctionEstimator{}, + sizingStrategy: DefaultSizingStrategy(), exprPaths: map[int64][]string{}, localVars: make(scopes), computedSizes: map[int64]SizeEstimate{}, @@ -157,6 +161,9 @@ func Cost(checked *ast.AST, estimator Estimator, opts ...Option) (CostEstimate, return CostEstimate{}, err } } + if c.sizingStrategy == nil { + c.sizingStrategy = DefaultSizingStrategy() + } return c.cost(checked.Expr()), nil } @@ -588,13 +595,7 @@ func (c *coster) functionCost(e ast.Expr, function, overloadID string, target *A return CallEstimate{CostEstimate: est.Add(argCost), ResultSize: est.ResultSize} } } - if c.sizingStrategy != nil { - if estimator, found := c.getSizingOverloadEstimators()[overloadID]; found { - if est := estimator(c.estimator, target, args); est != nil { - return CallEstimate{CostEstimate: est.Add(argCost), ResultSize: est.ResultSize} - } - } - } else if estimator, found := stdOverloadEstimators[overloadID]; found { + if estimator, found := c.getStandardOverloadEstimators()[overloadID]; found { if est := estimator(c.estimator, target, args); est != nil { return CallEstimate{CostEstimate: est.Add(argCost), ResultSize: est.ResultSize} } @@ -628,7 +629,10 @@ func (c *coster) getPath(e ast.Expr) []string { return nil } -func (c *coster) getSizingOverloadEstimators() map[string]FunctionEstimator { +func (c *coster) getStandardOverloadEstimators() map[string]FunctionEstimator { + if c.sizingStrategy == nil || c.sizingStrategy == defaultSizing { + return stdOverloadEstimators + } if c.sizingOverloadEstimators == nil { c.sizingOverloadEstimators = StandardOverloadEstimatorsWithOptions(c.sizingStrategy) } diff --git a/common/cost/model.go b/common/cost/model.go index 6dacee62..18d99b02 100644 --- a/common/cost/model.go +++ b/common/cost/model.go @@ -690,11 +690,14 @@ func (m OverloadModel) hasTarget() bool { // FunctionEstimator returns a FunctionEstimator implementing the cost model. func (m OverloadModel) FunctionEstimator() FunctionEstimator { - return m.FunctionEstimatorWithOptions(nil) + return m.FunctionEstimatorWithOptions(DefaultSizingStrategy()) } // FunctionEstimatorWithOptions returns a FunctionEstimator implementing the cost model with an optional SizingStrategy. func (m OverloadModel) FunctionEstimatorWithOptions(strategy SizingStrategy) FunctionEstimator { + if strategy == nil { + strategy = DefaultSizingStrategy() + } hasTarget := m.hasTarget() return func(estimator Estimator, target *AstNode, args []AstNode) *CallEstimate { if hasTarget && target == nil { @@ -719,11 +722,14 @@ func (m OverloadModel) FunctionEstimatorWithOptions(strategy SizingStrategy) Fun // FunctionTracker returns a FunctionTracker implementing the cost model. func (m OverloadModel) FunctionTracker() FunctionTracker { - return m.FunctionTrackerWithOptions(nil) + return m.FunctionTrackerWithOptions(DefaultSizingStrategy()) } // FunctionTrackerWithOptions returns a FunctionTracker implementing the cost model with an optional SizingStrategy. func (m OverloadModel) FunctionTrackerWithOptions(strategy SizingStrategy) FunctionTracker { + if strategy == nil { + strategy = DefaultSizingStrategy() + } isMember := m.hasTarget() return func(args []ref.Val, result ref.Val) *uint64 { ctx := &trackerEvalContext{ diff --git a/common/cost/standard.go b/common/cost/standard.go index d6bc6581..755ddc00 100644 --- a/common/cost/standard.go +++ b/common/cost/standard.go @@ -142,11 +142,14 @@ var StandardOverloadModels = []OverloadModel{ // StandardOverloadEstimators returns the map of FunctionEstimator instances for standard overloads. func StandardOverloadEstimators() map[string]FunctionEstimator { - return StandardOverloadEstimatorsWithOptions(nil) + return StandardOverloadEstimatorsWithOptions(DefaultSizingStrategy()) } // StandardOverloadEstimatorsWithOptions returns the map of FunctionEstimator instances for standard overloads with an optional SizingStrategy. func StandardOverloadEstimatorsWithOptions(strategy SizingStrategy) map[string]FunctionEstimator { + if strategy == nil { + strategy = DefaultSizingStrategy() + } estimators := make(map[string]FunctionEstimator, len(StandardOverloadModels)) for _, m := range StandardOverloadModels { estimators[m.ID] = m.FunctionEstimatorWithOptions(strategy) @@ -156,11 +159,14 @@ func StandardOverloadEstimatorsWithOptions(strategy SizingStrategy) map[string]F // StandardOverloadTrackers returns the map of FunctionTracker instances for standard overloads. func StandardOverloadTrackers() map[string]FunctionTracker { - return StandardOverloadTrackersWithOptions(nil) + return StandardOverloadTrackersWithOptions(DefaultSizingStrategy()) } // StandardOverloadTrackersWithOptions returns the map of FunctionTracker instances for standard overloads with an optional SizingStrategy. func StandardOverloadTrackersWithOptions(strategy SizingStrategy) map[string]FunctionTracker { + if strategy == nil { + strategy = DefaultSizingStrategy() + } trackers := make(map[string]FunctionTracker, len(StandardOverloadModels)) for _, m := range StandardOverloadModels { trackers[m.ID] = m.FunctionTrackerWithOptions(strategy) diff --git a/common/cost/tracker.go b/common/cost/tracker.go index c576f200..61bfd25f 100644 --- a/common/cost/tracker.go +++ b/common/cost/tracker.go @@ -84,6 +84,9 @@ func OverloadTracker(overloadID string, fnTracker FunctionTracker) TrackerOption // TrackerSizingStrategy configures a SizingStrategy for runtime cost tracking. func TrackerSizingStrategy(strategy SizingStrategy) TrackerOption { return func(tracker *Tracker) error { + if strategy == nil { + strategy = DefaultSizingStrategy() + } tracker.sizingStrategy = strategy return nil } @@ -117,6 +120,7 @@ func NewTracker(estimator ActualCostEstimator, opts ...TrackerOption) (*Tracker, tracker := &Tracker{ Estimator: estimator, overloadTrackers: map[string]FunctionTracker{}, + sizingStrategy: DefaultSizingStrategy(), presenceTestHasCost: true, } for _, opt := range opts { @@ -125,6 +129,9 @@ func NewTracker(estimator ActualCostEstimator, opts ...TrackerOption) (*Tracker, return nil, err } } + if tracker.sizingStrategy == nil { + tracker.sizingStrategy = DefaultSizingStrategy() + } return tracker, nil } @@ -224,7 +231,10 @@ func (c *Tracker) checkLimit() { } } -func (c *Tracker) getSizingOverloadTrackers() map[string]FunctionTracker { +func (c *Tracker) getStandardOverloadTrackers() map[string]FunctionTracker { + if c.sizingStrategy == nil || c.sizingStrategy == defaultSizing { + return stdOverloadTrackers + } if c.sizingOverloadTrackers == nil { c.sizingOverloadTrackers = StandardOverloadTrackersWithOptions(c.sizingStrategy) } @@ -250,15 +260,7 @@ func (c *Tracker) CostCall(call Call, args []ref.Val, result ref.Val) uint64 { return total } } - if c.sizingStrategy != nil { - if tracker, found := c.getSizingOverloadTrackers()[call.OverloadID()]; found { - callCost := tracker(args, result) - if callCost != nil { - total = SafeAdd(total, *callCost) - return total - } - } - } else if tracker, found := stdOverloadTrackers[call.OverloadID()]; found { + if tracker, found := c.getStandardOverloadTrackers()[call.OverloadID()]; found { callCost := tracker(args, result) if callCost != nil { total = SafeAdd(total, *callCost) diff --git a/interpreter/interpretable.go b/interpreter/interpretable.go index 4713897d..f177ade4 100644 --- a/interpreter/interpretable.go +++ b/interpreter/interpretable.go @@ -444,7 +444,7 @@ func (eq *evalEq) Exec(frame *ExecutionFrame) ref.Val { if types.IsError(lVal) { // To preserve legacy cost tracking behavior for == // track this cost, but it will be removed in the future. - trackCostEvalBinary(frame, eq.id, eq, lVal, nil, lVal) + trackCostEvalBinary(frame, eq.id, eq, lVal, types.UnknownType, lVal) return lVal } rVal := eq.rhs.Exec(frame) @@ -503,7 +503,7 @@ func (ne *evalNe) Exec(frame *ExecutionFrame) ref.Val { if types.IsError(lVal) { // To preserve legacy cost tracking behavior for !=, // track this cost, but it will be removed in the future. - trackCostEvalBinary(frame, ne.id, ne, lVal, nil, lVal) + trackCostEvalBinary(frame, ne.id, ne, lVal, types.UnknownType, lVal) return lVal } rVal := ne.rhs.Exec(frame) From 781243e5937d838f5f0b20e5ebaf777d6a884826 Mon Sep 17 00:00:00 2001 From: TristonianJones Date: Thu, 10 Sep 2026 14:13:03 -0700 Subject: [PATCH 4/6] Additional simplifications to the cost model --- common/cost/cost.go | 33 +++++- common/cost/cost_test.go | 112 +++++++++++++++++++- common/cost/estimator.go | 27 +++-- common/cost/model.go | 177 ++++++++++++++++++++++++++++++++ common/cost/model_test.go | 197 +++++++++++++++++++++++++++++++++--- common/cost/tracker.go | 13 ++- common/cost/tracker_test.go | 48 ++++++++- ext/costs.go | 2 +- 8 files changed, 568 insertions(+), 41 deletions(-) diff --git a/common/cost/cost.go b/common/cost/cost.go index 15cba95f..a7fda4b3 100644 --- a/common/cost/cost.go +++ b/common/cost/cost.go @@ -90,6 +90,14 @@ func SafeAdd(x, y uint64, rest ...uint64) uint64 { return sum } +// SafeSubtract returns the difference of x - y, saturating at zero. +func SafeSubtract(x, y uint64) uint64 { + if x > y { + return x - y + } + return 0 +} + // SafeMultiply returns the product of the input values, saturating at math.MaxUint64. func SafeMultiply(x, y uint64) uint64 { if y != 0 && x > math.MaxUint64/y { @@ -153,8 +161,8 @@ func RangedSizeEstimate(min, max uint64) SizeEstimate { return SizeEstimate{Min: min, Max: max} } -// AtLeastOne returns a size estimate with min and max guaranteed to be at least 1. -func AtLeastOne(size SizeEstimate) SizeEstimate { +// AtLeastOneSize returns a size estimate with min and max guaranteed to be at least 1. +func AtLeastOneSize(size SizeEstimate) SizeEstimate { if size.Min == 0 { size.Min = 1 } @@ -189,6 +197,17 @@ func (se SizeEstimate) Add(sizeEstimate SizeEstimate) SizeEstimate { return res } +// Subtract subtracts another SizeEstimate and returns the difference, saturating at zero. +func (se SizeEstimate) Subtract(sizeEstimate SizeEstimate) SizeEstimate { + res := SizeEstimate{ + Min: SafeSubtract(se.Min, sizeEstimate.Max), + Max: SafeSubtract(se.Max, sizeEstimate.Min), + } + res.Key = mergeSizeEstimatePtr(se.Key, sizeEstimate.Key) + res.Elem = mergeSizeEstimatePtr(se.Elem, sizeEstimate.Elem) + return res +} + // Multiply multiplies by another SizeEstimate and returns the product. // If multiply would result in an uint64 overflow, the result is math.MaxUint64. func (se SizeEstimate) Multiply(sizeEstimate SizeEstimate) SizeEstimate { @@ -326,6 +345,9 @@ func ActualSize(value ref.Val) uint64 { // EstimateSize returns a SizeEstimate for the given node from its computed size, estimator, or unknown. func EstimateSize(estimator Estimator, node AstNode) SizeEstimate { + if node == nil { + return UnknownSizeEstimate() + } if l := node.ComputedSize(); l != nil { return *l } @@ -359,10 +381,13 @@ func EstimateListAlloc(sz SizeEstimate, costFactor float64) (CostEstimate, *Size // NodeAsUintValue returns the value of a literal int node as a uint64, or the default value if the // node is not a non-negative int literal. func NodeAsUintValue(node AstNode, defaultVal uint64) uint64 { - if node.Expr().Kind() != ast.LiteralKind { + if node == nil || node.Expr() == nil || node.Expr().Kind() != ast.LiteralKind { return defaultVal } lit := node.Expr().AsLiteral() + if lit.Type() == types.UintType { + return uint64(lit.(types.Uint)) + } if lit.Type() != types.IntType { return defaultVal } @@ -370,5 +395,5 @@ func NodeAsUintValue(node AstNode, defaultVal uint64) uint64 { if val < types.IntZero { return 0 } - return uint64(lit.(types.Int)) + return uint64(val) } diff --git a/common/cost/cost_test.go b/common/cost/cost_test.go index 5e18dad6..2922a28b 100644 --- a/common/cost/cost_test.go +++ b/common/cost/cost_test.go @@ -17,8 +17,32 @@ package cost import ( "math" "testing" + + "cel.dev/cel-go/common/ast" + "cel.dev/cel-go/common/types" ) +func TestSafeSubtract(t *testing.T) { + tests := []struct { + name string + x, y uint64 + want uint64 + }{ + {name: "zero", x: 0, y: 0, want: 0}, + {name: "simple", x: 5, y: 3, want: 2}, + {name: "underflow to zero", x: 3, y: 5, want: 0}, + {name: "max minus zero", x: math.MaxUint64, y: 0, want: math.MaxUint64}, + {name: "max minus max", x: math.MaxUint64, y: math.MaxUint64, want: 0}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := SafeSubtract(tc.x, tc.y); got != tc.want { + t.Errorf("SafeSubtract(%d, %d) got %d, want %d", tc.x, tc.y, got, tc.want) + } + }) + } +} + func TestSafeAdd(t *testing.T) { tests := []struct { name string @@ -144,7 +168,7 @@ func TestSizeEstimate(t *testing.T) { if got := RangedSizeEstimate(3, 8); got.Min != 3 || got.Max != 8 { t.Errorf("RangedSizeEstimate(3, 8) = %v, want {3, 8}", got) } - if got := AtLeastOne(FixedSizeEstimate(0)); got.Min != 1 || got.Max != 1 { + if got := AtLeastOneSize(FixedSizeEstimate(0)); got.Min != 1 || got.Max != 1 { t.Errorf("AtLeastOne(0) = %v, want {1, 1}", got) } } @@ -192,3 +216,89 @@ func TestExtCostHelpers(t *testing.T) { t.Errorf("NewCallEstimate = %v, want CostEstimate=%v ResultSize=%v", callEst, costEst, resSz) } } + +func TestSizeEstimate_Subtract(t *testing.T) { + s1 := RangedSizeEstimate(5, 15) + s2 := RangedSizeEstimate(2, 4) + got := s1.Subtract(s2) + if got.Min != 1 || got.Max != 13 { + t.Errorf("s1.Subtract(s2) = %v, want {1, 13}", got) + } + + // Underflow cases + s3 := RangedSizeEstimate(2, 4) + s4 := RangedSizeEstimate(5, 10) + got2 := s3.Subtract(s4) + if got2.Min != 0 || got2.Max != 0 { + t.Errorf("s3.Subtract(s4) = %v, want {0, 0}", got2) + } +} + +func TestEstimateSize(t *testing.T) { + // Nil node + if sz := EstimateSize(nil, nil); sz != UnknownSizeEstimate() { + t.Errorf("EstimateSize(nil, nil) = %v, want unknown", sz) + } + + // Node with computed size + compSz := FixedSizeEstimate(42) + nodeWithComp := NewAstNode(nil, nil, types.IntType, &compSz) + if sz := EstimateSize(nil, nodeWithComp); sz != compSz { + t.Errorf("EstimateSize(comp) = %v, want %v", sz, compSz) + } + + // Node with estimator + nodeWithoutComp := NewAstNode(nil, []string{"foo"}, types.IntType, nil) + est := testHintsEstimator{hints: map[string]uint64{"foo": 100}} + if sz := EstimateSize(est, nodeWithoutComp); sz != FixedSizeEstimate(100) { + t.Errorf("EstimateSize(est) = %v, want 100", sz) + } + + // Node with estimator returning nil + if sz := EstimateSize(est, NewAstNode(nil, []string{"bar"}, types.IntType, nil)); sz != UnknownSizeEstimate() { + t.Errorf("EstimateSize(unknown) = %v, want unknown", sz) + } +} + +func TestNodeAsUintValue(t *testing.T) { + // Nil node + if val := NodeAsUintValue(nil, 99); val != 99 { + t.Errorf("NodeAsUintValue(nil) = %d, want 99", val) + } + + // Non-literal node + fac := ast.NewExprFactory() + identExpr := fac.NewIdent(1, "x") + identNode := NewAstNode(identExpr, nil, types.IntType, nil) + if val := NodeAsUintValue(identNode, 99); val != 99 { + t.Errorf("NodeAsUintValue(ident) = %d, want 99", val) + } + + // Non-int literal (string) + strExpr := fac.NewLiteral(2, types.String("hello")) + strNode := NewAstNode(strExpr, nil, types.StringType, nil) + if val := NodeAsUintValue(strNode, 99); val != 99 { + t.Errorf("NodeAsUintValue(str) = %d, want 99", val) + } + + // Positive int literal + intExpr := fac.NewLiteral(3, types.Int(42)) + intNode := NewAstNode(intExpr, nil, types.IntType, nil) + if val := NodeAsUintValue(intNode, 99); val != 42 { + t.Errorf("NodeAsUintValue(int 42) = %d, want 42", val) + } + + // Negative int literal (saturates at 0) + negExpr := fac.NewLiteral(4, types.Int(-5)) + negNode := NewAstNode(negExpr, nil, types.IntType, nil) + if val := NodeAsUintValue(negNode, 99); val != 0 { + t.Errorf("NodeAsUintValue(int -5) = %d, want 0", val) + } + + // Uint literal + uintExpr := fac.NewLiteral(5, types.Uint(100)) + uintNode := NewAstNode(uintExpr, nil, types.UintType, nil) + if val := NodeAsUintValue(uintNode, 99); val != 100 { + t.Errorf("NodeAsUintValue(uint 100) = %d, want 100", val) + } +} diff --git a/common/cost/estimator.go b/common/cost/estimator.go index 7260ef19..1e974c7d 100644 --- a/common/cost/estimator.go +++ b/common/cost/estimator.go @@ -149,7 +149,6 @@ func Cost(checked *ast.AST, estimator Estimator, opts ...Option) (CostEstimate, checkedAST: checked, estimator: estimator, overloadEstimators: map[string]FunctionEstimator{}, - sizingStrategy: DefaultSizingStrategy(), exprPaths: map[int64][]string{}, localVars: make(scopes), computedSizes: map[int64]SizeEstimate{}, @@ -164,6 +163,9 @@ func Cost(checked *ast.AST, estimator Estimator, opts ...Option) (CostEstimate, if c.sizingStrategy == nil { c.sizingStrategy = DefaultSizingStrategy() } + if c.sizingStrategy != defaultSizing { + c.sizingOverloadEstimators = StandardOverloadEstimatorsWithOptions(c.sizingStrategy) + } return c.cost(checked.Expr()), nil } @@ -630,13 +632,10 @@ func (c *coster) getPath(e ast.Expr) []string { } func (c *coster) getStandardOverloadEstimators() map[string]FunctionEstimator { - if c.sizingStrategy == nil || c.sizingStrategy == defaultSizing { - return stdOverloadEstimators - } - if c.sizingOverloadEstimators == nil { - c.sizingOverloadEstimators = StandardOverloadEstimatorsWithOptions(c.sizingStrategy) + if c.sizingOverloadEstimators != nil { + return c.sizingOverloadEstimators } - return c.sizingOverloadEstimators + return stdOverloadEstimators } // addPath associates an expression ID with its path. @@ -731,6 +730,20 @@ func (e *estimatorContext) ArgType(index int) (*types.Type, bool) { return nil, false } +func (e *estimatorContext) ArgValue(index int, defaultVal uint64) uint64 { + if index < len(e.args) && e.args[index] != nil { + return NodeAsUintValue(e.args[index], defaultVal) + } + return defaultVal +} + +func (e *estimatorContext) TargetValue(defaultVal uint64) uint64 { + if e.target != nil && (*e.target) != nil { + return NodeAsUintValue(*e.target, defaultVal) + } + return defaultVal +} + func (e *estimatorContext) Size(node AstNode) SizeEstimate { if node == nil { return UnknownSizeEstimate() diff --git a/common/cost/model.go b/common/cost/model.go index 18d99b02..2ab7a939 100644 --- a/common/cost/model.go +++ b/common/cost/model.go @@ -48,6 +48,12 @@ type EstimateContext interface { // Size returns the estimated size based on the SizingStrategy configured on the Estimator. Size(node AstNode) SizeEstimate + + // ArgValue returns the literal/constant uint64 value of the argument at index, or defaultVal. + ArgValue(index int, defaultVal uint64) uint64 + + // TargetValue returns the literal/constant uint64 value of the receiver/target, or defaultVal. + TargetValue(defaultVal uint64) uint64 } // TrackContext provides value size and argument evaluation during runtime cost tracking. @@ -68,6 +74,12 @@ type TrackContext interface { // Size returns the actual runtime size of a value. Size(value ref.Val) uint64 + + // ArgValue returns the uint64 value of the argument at index, or defaultVal. + ArgValue(index int, defaultVal uint64) uint64 + + // TargetValue returns the uint64 value of the receiver/target, or defaultVal. + TargetValue(defaultVal uint64) uint64 } // QuantityExpr represents a computable size or cost equation. @@ -145,6 +157,28 @@ func ArgKey(index int) QuantityExpr { return KeyOf(Arg(index)) } +// intArgExpr represents the integer value of an argument. +type intArgExpr struct { + index int + defaultVal uint64 +} + +func (a intArgExpr) estimate(ctx EstimateContext) SizeEstimate { + return FixedSizeEstimate(ctx.ArgValue(a.index, a.defaultVal)) +} + +func (a intArgExpr) track(ctx TrackContext) uint64 { + return ctx.ArgValue(a.index, a.defaultVal) +} + +func (intArgExpr) hasTarget() bool { return false } + +// IntArg creates an expression referencing the integer value of the argument at the given index, +// falling back to defaultVal if the argument is not a non-negative integer. +func IntArg(index int, defaultVal uint64) QuantityExpr { + return intArgExpr{index: index, defaultVal: defaultVal} +} + // targetExpr represents the size of the receiver/target object. type targetExpr struct{} @@ -202,6 +236,27 @@ func TargetKey() QuantityExpr { return KeyOf(Target()) } +// intTargetExpr represents the integer value of the receiver/target object. +type intTargetExpr struct { + defaultVal uint64 +} + +func (t intTargetExpr) estimate(ctx EstimateContext) SizeEstimate { + return FixedSizeEstimate(ctx.TargetValue(t.defaultVal)) +} + +func (t intTargetExpr) track(ctx TrackContext) uint64 { + return ctx.TargetValue(t.defaultVal) +} + +func (intTargetExpr) hasTarget() bool { return true } + +// IntTarget creates an expression referencing the integer value of the receiver/target object, +// falling back to defaultVal if the target is not a non-negative integer. +func IntTarget(defaultVal uint64) QuantityExpr { + return intTargetExpr{defaultVal: defaultVal} +} + // keyExpr represents the key size of another quantity expression. type keyExpr struct { expr QuantityExpr @@ -282,6 +337,30 @@ func Sum(terms ...QuantityExpr) QuantityExpr { return addExpr{terms: terms} } +// subExpr represents the subtraction of two quantity expressions. +type subExpr struct { + lhs, rhs QuantityExpr +} + +func (s subExpr) estimate(ctx EstimateContext) SizeEstimate { + lhsVal := s.lhs.estimate(ctx) + rhsVal := s.rhs.estimate(ctx) + return lhsVal.Subtract(rhsVal) +} + +func (s subExpr) track(ctx TrackContext) uint64 { + return SafeSubtract(s.lhs.track(ctx), s.rhs.track(ctx)) +} + +func (s subExpr) hasTarget() bool { + return hasTarget(s.lhs) || hasTarget(s.rhs) +} + +// Sub creates an expression representing lhs - rhs, saturated at zero. +func Sub(lhs, rhs QuantityExpr) QuantityExpr { + return subExpr{lhs: lhs, rhs: rhs} +} + // mulExpr represents the product of multiple quantity expressions. type mulExpr struct { terms []QuantityExpr @@ -586,6 +665,50 @@ func AtMost(maxExpr QuantityExpr) QuantityExpr { return Ranged(Const(0), maxExpr) } +// atLeastOneExpr represents a quantity guaranteed to be at least 1. +type atLeastOneExpr struct { + expr QuantityExpr +} + +func (a atLeastOneExpr) estimate(ctx EstimateContext) SizeEstimate { + return AtLeastOneSize(a.expr.estimate(ctx)) +} + +func (a atLeastOneExpr) track(ctx TrackContext) uint64 { + val := a.expr.track(ctx) + if val == 0 { + return 1 + } + return val +} + +func (a atLeastOneExpr) hasTarget() bool { + return hasTarget(a.expr) +} + +// AtLeastOneQuantity creates an expression ensuring the quantity is at least 1. +func AtLeastOneQuantity(expr QuantityExpr) QuantityExpr { + return atLeastOneExpr{expr: expr} +} + +// StringScan creates an expression representing the cost of scanning a string. +func StringScan(expr QuantityExpr) QuantityExpr { + return Scale(expr, StringTraversalCostFactor) +} + +// ListAlloc creates an expression representing list allocation cost with base cost and scaled element count. +func ListAlloc(elemCount QuantityExpr, costFactor float64) QuantityExpr { + return Sum(Const(ListCreateBaseCost), Scale(elemCount, costFactor)) +} + +// Traversal creates an expression representing the traversal and optional allocation cost over an expression. +func Traversal(expr QuantityExpr, costFactor float64, allocCost uint64) QuantityExpr { + if allocCost == 0 { + return Scale(expr, costFactor) + } + return Sum(Const(allocCost), Scale(expr, costFactor)) +} + // listExpr represents a list size estimate composed of length and element size expressions. type listExpr struct { lenExpr QuantityExpr @@ -815,6 +938,22 @@ func (e *estimatorEvalContext) ArgType(index int) (*types.Type, bool) { return nil, false } +// ArgValue returns the uint64 value of the argument at index, or defaultVal if not a literal int/uint. +func (e *estimatorEvalContext) ArgValue(index int, defaultVal uint64) uint64 { + if index < len(e.args) && e.args[index] != nil { + return NodeAsUintValue(e.args[index], defaultVal) + } + return defaultVal +} + +// TargetValue returns the uint64 value of the receiver/target, or defaultVal if not a literal int/uint. +func (e *estimatorEvalContext) TargetValue(defaultVal uint64) uint64 { + if e.target != nil && (*e.target) != nil { + return NodeAsUintValue(*e.target, defaultVal) + } + return defaultVal +} + // trackerEvalContext provides evaluation context for runtime cost tracking. type trackerEvalContext struct { estimator ActualCostEstimator @@ -824,6 +963,44 @@ type trackerEvalContext struct { isMember bool } +// valueAsUint returns the non-negative uint64 value of ref.Val (int or uint), or defaultVal. +func valueAsUint(val ref.Val, defaultVal uint64) uint64 { + if val == nil { + return defaultVal + } + switch v := val.(type) { + case types.Int: + if v < 0 { + return 0 + } + return uint64(v) + case types.Uint: + return uint64(v) + default: + return defaultVal + } +} + +// ArgValue returns the uint64 value of the argument at index, or defaultVal. +func (t *trackerEvalContext) ArgValue(index int, defaultVal uint64) uint64 { + idx := index + if t.isMember { + idx = index + 1 + } + if idx < len(t.args) { + return valueAsUint(t.args[idx], defaultVal) + } + return defaultVal +} + +// TargetValue returns the uint64 value of the receiver/target object, or defaultVal. +func (t *trackerEvalContext) TargetValue(defaultVal uint64) uint64 { + if t.isMember && len(t.args) > 0 { + return valueAsUint(t.args[0], defaultVal) + } + return defaultVal +} + // TargetType returns the type of the receiver/target object. func (t *trackerEvalContext) TargetType() (*types.Type, bool) { if t.isMember && len(t.args) > 0 && t.args[0] != nil { diff --git a/common/cost/model_test.go b/common/cost/model_test.go index 5ba5db66..0b063ef2 100644 --- a/common/cost/model_test.go +++ b/common/cost/model_test.go @@ -23,13 +23,29 @@ import ( ) type testEvalContext struct { - estimator Estimator - strategy SizingStrategy - args []SizeEstimate - receiver *SizeEstimate - result *SizeEstimate - targetType *types.Type - argTypes []*types.Type + estimator Estimator + strategy SizingStrategy + args []SizeEstimate + argValues []uint64 + receiver *SizeEstimate + receiverValue *uint64 + result *SizeEstimate + targetType *types.Type + argTypes []*types.Type +} + +func (t *testEvalContext) ArgValue(index int, defaultVal uint64) uint64 { + if index < len(t.argValues) { + return t.argValues[index] + } + return defaultVal +} + +func (t *testEvalContext) TargetValue(defaultVal uint64) uint64 { + if t.receiverValue != nil { + return *t.receiverValue + } + return defaultVal } func (t *testEvalContext) Arg(index int) (SizeEstimate, bool) { @@ -105,14 +121,17 @@ func TestQuantityExprs_Estimate(t *testing.T) { arg1 := MapSizeEstimate(RangedSizeEstimate(3, 5), key1, elem1) rcv := MapSizeEstimate(RangedSizeEstimate(4, 8), rcvKey, rcvElem) + receiverVal := uint64(20) ctx := &testEvalContext{ args: []SizeEstimate{ arg0, arg1, }, - receiver: &rcv, - targetType: types.NewListType(types.StringType), - argTypes: []*types.Type{types.NewListType(types.IntType)}, + argValues: []uint64{3, 10}, + receiver: &rcv, + receiverValue: &receiverVal, + targetType: types.NewListType(types.StringType), + argTypes: []*types.Type{types.NewListType(types.IntType)}, } tests := []struct { @@ -130,6 +149,16 @@ func TestQuantityExprs_Estimate(t *testing.T) { expr: Arg(0), expected: arg0, }, + { + name: "int_arg_quantity", + expr: IntArg(0, 99), + expected: FixedSizeEstimate(3), + }, + { + name: "int_arg_default_quantity", + expr: IntArg(5, 99), + expected: FixedSizeEstimate(99), + }, { name: "arg_element_quantity", expr: ArgElem(0), @@ -145,6 +174,11 @@ func TestQuantityExprs_Estimate(t *testing.T) { expr: Target(), expected: rcv, }, + { + name: "int_target_quantity", + expr: IntTarget(99), + expected: FixedSizeEstimate(20), + }, { name: "target_element_quantity", expr: TargetElem(), @@ -160,6 +194,11 @@ func TestQuantityExprs_Estimate(t *testing.T) { expr: Sum(Arg(0), Arg(1)), expected: arg0.Add(arg1), }, + { + name: "sub_quantity", + expr: Sub(Arg(0), Arg(1)), + expected: arg0.Subtract(arg1), + }, { name: "mul_quantity", expr: Mul(Arg(0), Arg(1)), @@ -170,6 +209,31 @@ func TestQuantityExprs_Estimate(t *testing.T) { expr: Scale(Arg(0), 0.5), expected: ListSizeEstimate(RangedSizeEstimate(1, 5), elem0), }, + { + name: "string_scan_quantity", + expr: StringScan(Const(20)), + expected: FixedSizeEstimate(2), + }, + { + name: "list_alloc_quantity", + expr: ListAlloc(Const(10), 0.5), + expected: FixedSizeEstimate(ListCreateBaseCost + 5), + }, + { + name: "traversal_with_alloc_quantity", + expr: Traversal(Const(20), 0.5, 10), + expected: FixedSizeEstimate(20), + }, + { + name: "traversal_without_alloc_quantity", + expr: Traversal(Const(20), 0.5, 0), + expected: FixedSizeEstimate(10), + }, + { + name: "at_least_one_quantity", + expr: AtLeastOneQuantity(Const(0)), + expected: FixedSizeEstimate(1), + }, { name: "square_quantity", expr: Square(Arg(1)), @@ -263,9 +327,12 @@ func TestQuantityExprs_Estimate(t *testing.T) { } func TestQuantityExprs_Track(t *testing.T) { + rcvTrackVal := uint64(20) trackCtx := &testTrackContext{ - args: []uint64{10, 5}, - receiver: 8, + args: []uint64{10, 5}, + argValues: []uint64{3, 10}, + receiver: 8, + receiverValue: &rcvTrackVal, } scalarTests := []struct { @@ -275,10 +342,21 @@ func TestQuantityExprs_Track(t *testing.T) { }{ {name: "const_scalar", expr: Const(42), expected: 42}, {name: "arg_scalar", expr: Arg(0), expected: 10}, + {name: "int_arg_scalar", expr: IntArg(0, 99), expected: 3}, + {name: "int_arg_default_scalar", expr: IntArg(5, 99), expected: 99}, {name: "target_scalar", expr: Target(), expected: 8}, + {name: "int_target_scalar", expr: IntTarget(99), expected: 20}, {name: "sum_scalar", expr: Sum(Arg(0), Arg(1)), expected: 15}, + {name: "sub_scalar", expr: Sub(Arg(0), Arg(1)), expected: 5}, + {name: "sub_scalar_saturates", expr: Sub(Arg(1), Arg(0)), expected: 0}, {name: "mul_scalar", expr: Mul(Arg(0), Arg(1)), expected: 50}, {name: "scale_scalar", expr: Scale(Arg(0), 1.5), expected: 15}, + {name: "string_scan_scalar", expr: StringScan(Const(20)), expected: 2}, + {name: "list_alloc_scalar", expr: ListAlloc(Const(10), 0.5), expected: ListCreateBaseCost + 5}, + {name: "traversal_with_alloc_scalar", expr: Traversal(Const(20), 0.5, 10), expected: 20}, + {name: "traversal_without_alloc_scalar", expr: Traversal(Const(20), 0.5, 0), expected: 10}, + {name: "at_least_one_scalar_zero", expr: AtLeastOneQuantity(Const(0)), expected: 1}, + {name: "at_least_one_scalar_nonzero", expr: AtLeastOneQuantity(Const(5)), expected: 5}, {name: "square_scalar", expr: Square(Arg(1)), expected: 25}, {name: "min_scalar", expr: Min(Arg(0), Arg(1)), expected: 5}, {name: "max_scalar", expr: Max(Arg(0), Arg(1)), expected: 10}, @@ -296,10 +374,26 @@ func TestQuantityExprs_Track(t *testing.T) { } type testTrackContext struct { - args []uint64 - receiver uint64 - result uint64 - estimator ActualCostEstimator + args []uint64 + argValues []uint64 + receiver uint64 + receiverValue *uint64 + result uint64 + estimator ActualCostEstimator +} + +func (t *testTrackContext) ArgValue(index int, defaultVal uint64) uint64 { + if index < len(t.argValues) { + return t.argValues[index] + } + return defaultVal +} + +func (t *testTrackContext) TargetValue(defaultVal uint64) uint64 { + if t.receiverValue != nil { + return *t.receiverValue + } + return defaultVal } func (t *testTrackContext) Arg(index int) uint64 { @@ -350,7 +444,9 @@ func TestModel_HasTargetInspection(t *testing.T) { {name: "arg_expr", expr: Arg(0), wantTarget: false}, {name: "arg_elem_expr", expr: ArgElem(0), wantTarget: false}, {name: "arg_key_expr", expr: ArgKey(0), wantTarget: false}, + {name: "int_arg_expr", expr: IntArg(0, 10), wantTarget: false}, {name: "target_expr", expr: Target(), wantTarget: true}, + {name: "int_target_expr", expr: IntTarget(10), wantTarget: true}, {name: "target_elem_expr", expr: TargetElem(), wantTarget: true}, {name: "target_key_expr", expr: TargetKey(), wantTarget: true}, {name: "result_expr", expr: Result(), wantTarget: false}, @@ -358,6 +454,17 @@ func TestModel_HasTargetInspection(t *testing.T) { {name: "elem_of_arg", expr: ElemOf(Arg(0)), wantTarget: false}, {name: "key_of_target", expr: KeyOf(Target()), wantTarget: true}, {name: "key_of_arg", expr: KeyOf(Arg(0)), wantTarget: false}, + {name: "sub_with_target_lhs", expr: Sub(Target(), Arg(0)), wantTarget: true}, + {name: "sub_with_target_rhs", expr: Sub(Arg(0), Target()), wantTarget: true}, + {name: "sub_without_target", expr: Sub(Arg(0), Arg(1)), wantTarget: false}, + {name: "at_least_one_with_target", expr: AtLeastOneQuantity(Target()), wantTarget: true}, + {name: "at_least_one_without_target", expr: AtLeastOneQuantity(Arg(0)), wantTarget: false}, + {name: "string_scan_with_target", expr: StringScan(Target()), wantTarget: true}, + {name: "string_scan_without_target", expr: StringScan(Arg(0)), wantTarget: false}, + {name: "list_alloc_with_target", expr: ListAlloc(Target(), 1.0), wantTarget: true}, + {name: "list_alloc_without_target", expr: ListAlloc(Arg(0), 1.0), wantTarget: false}, + {name: "traversal_with_target", expr: Traversal(Target(), 1.0, 10), wantTarget: true}, + {name: "traversal_without_target", expr: Traversal(Arg(0), 1.0, 10), wantTarget: false}, {name: "min_with_target", expr: Min(Arg(0), Target()), wantTarget: true}, {name: "min_without_target", expr: Min(Arg(0), Arg(1)), wantTarget: false}, {name: "max_with_target", expr: Max(Target(), Arg(0)), wantTarget: true}, @@ -539,6 +646,26 @@ func TestModel_MissingContextFallbacks(t *testing.T) { } }, }, + { + name: "missing_int_arg", + expr: IntArg(0, 42), + wantTrack: 42, + checkEst: func(t *testing.T, sz SizeEstimate) { + if sz != FixedSizeEstimate(42) { + t.Errorf("got %v, want 42", sz) + } + }, + }, + { + name: "missing_int_target", + expr: IntTarget(42), + wantTrack: 42, + checkEst: func(t *testing.T, sz SizeEstimate) { + if sz != FixedSizeEstimate(42) { + t.Errorf("got %v, want 42", sz) + } + }, + }, { name: "missing_result", expr: Result(), @@ -891,6 +1018,25 @@ func TestModel_TrackerEvalContextMethods(t *testing.T) { } }, }, + { + name: "target_value", + check: func(t *testing.T) { + if val := tCtx.TargetValue(99); val != 99 { + t.Errorf("TargetValue(99) for string target = %d, want 99", val) + } + }, + }, + { + name: "arg_value", + check: func(t *testing.T) { + if val := tCtx.ArgValue(0, 99); val != 100 { + t.Errorf("ArgValue(0, 99) = %d, want 100", val) + } + if val := tCtx.ArgValue(5, 99); val != 99 { + t.Errorf("ArgValue(5, 99) = %d, want 99", val) + } + }, + }, } for _, tc := range tests { @@ -1014,6 +1160,25 @@ func TestModel_EstimatorEvalContextMethods(t *testing.T) { } }, }, + { + name: "target_value", + check: func(t *testing.T) { + if val := eCtx.TargetValue(99); val != 99 { + t.Errorf("TargetValue(99) = %d, want 99", val) + } + }, + }, + { + name: "arg_value", + check: func(t *testing.T) { + if val := eCtx.ArgValue(0, 99); val != 99 { + t.Errorf("ArgValue(0, 99) = %d, want 99", val) + } + if val := eCtx.ArgValue(5, 99); val != 99 { + t.Errorf("ArgValue(5, 99) = %d, want 99", val) + } + }, + }, { name: "size_nil_node", check: func(t *testing.T) { diff --git a/common/cost/tracker.go b/common/cost/tracker.go index 61bfd25f..59e46bf1 100644 --- a/common/cost/tracker.go +++ b/common/cost/tracker.go @@ -120,7 +120,6 @@ func NewTracker(estimator ActualCostEstimator, opts ...TrackerOption) (*Tracker, tracker := &Tracker{ Estimator: estimator, overloadTrackers: map[string]FunctionTracker{}, - sizingStrategy: DefaultSizingStrategy(), presenceTestHasCost: true, } for _, opt := range opts { @@ -132,6 +131,9 @@ func NewTracker(estimator ActualCostEstimator, opts ...TrackerOption) (*Tracker, if tracker.sizingStrategy == nil { tracker.sizingStrategy = DefaultSizingStrategy() } + if tracker.sizingStrategy != defaultSizing { + tracker.sizingOverloadTrackers = StandardOverloadTrackersWithOptions(tracker.sizingStrategy) + } return tracker, nil } @@ -232,13 +234,10 @@ func (c *Tracker) checkLimit() { } func (c *Tracker) getStandardOverloadTrackers() map[string]FunctionTracker { - if c.sizingStrategy == nil || c.sizingStrategy == defaultSizing { - return stdOverloadTrackers - } - if c.sizingOverloadTrackers == nil { - c.sizingOverloadTrackers = StandardOverloadTrackersWithOptions(c.sizingStrategy) + if c.sizingOverloadTrackers != nil { + return c.sizingOverloadTrackers } - return c.sizingOverloadTrackers + return stdOverloadTrackers } // CostCall calculates the runtime cost for a function call. diff --git a/common/cost/tracker_test.go b/common/cost/tracker_test.go index dac6fcc9..dace09cb 100644 --- a/common/cost/tracker_test.go +++ b/common/cost/tracker_test.go @@ -1,4 +1,4 @@ -// Copyright 2022 Google LLC +// Copyright 2026 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -20,6 +20,7 @@ import ( "math/rand" "reflect" "strings" + "sync" "testing" "time" @@ -1166,8 +1167,8 @@ func TestRuntimeCost(t *testing.T) { }, { name: "two-var all: map literal early return false", - expr: `{'hello': 'world', 'taco': 'taco'}.all(k, v, k != v) == false`, - want: 44, + expr: `{'hello': 'world'}.all(k, v, k != v) == false`, + want: 38, }, { name: "two-var exists: list literal", @@ -1176,8 +1177,8 @@ func TestRuntimeCost(t *testing.T) { }, { name: "two-var exists: map literal", - expr: `{"a": 1, "b": 2}.exists(k, v, k == "a" && v == 1)`, - want: 42, + expr: `{"a": 1}.exists(k, v, k == "a" && v == 1)`, + want: 39, }, { name: "two-var existsOne: list variable", @@ -1418,3 +1419,40 @@ func newStandardInterpreter(t testing.TB, } return interpreter.NewInterpreter(disp, container, provider, adapter, resolver) } + +type testConcurrentSizingStrategy struct{} + +func (testConcurrentSizingStrategy) EstimateSize(ctx cost.EstimateContext, node cost.AstNode) (cost.SizeEstimate, bool) { + return cost.FixedSizeEstimate(10), true +} + +func (testConcurrentSizingStrategy) TrackSize(ctx cost.TrackContext, value ref.Val) (uint64, bool) { + return 10, true +} + +func TestTracker_ConcurrentCloneRace(t *testing.T) { + tracker, err := cost.NewTracker(nil, cost.TrackerSizingStrategy(testConcurrentSizingStrategy{})) + if err != nil { + t.Fatalf("NewTracker() failed: %v", err) + } + + var wg sync.WaitGroup + for i := 0; i < 50; i++ { + wg.Add(1) + go func() { + defer wg.Done() + clone, err := tracker.Clone() + if err != nil { + t.Errorf("tracker.Clone() failed: %v", err) + return + } + clone.CostCall(testCall{function: "startsWith", overloadID: overloads.StartsWithString}, []ref.Val{types.String("hello"), types.String("h")}, types.True) + clone.CostCall(testCall{function: "_==_", overloadID: overloads.Equals}, []ref.Val{types.String("a"), types.String("b")}, types.False) + clone.CreateList(1, nil) + if clone.ActualCost() == 0 { + t.Errorf("clone.ActualCost() should be non-zero") + } + }() + } + wg.Wait() +} diff --git a/ext/costs.go b/ext/costs.go index add7ede2..1b9f6c25 100644 --- a/ext/costs.go +++ b/ext/costs.go @@ -66,5 +66,5 @@ func fixedSizeEstimate(val uint64) cost.SizeEstimate { } func atLeastOne(size cost.SizeEstimate) cost.SizeEstimate { - return cost.AtLeastOne(size) + return cost.AtLeastOneSize(size) } From 0d3022c18d1554079ec44c6db93976232bce0deb Mon Sep 17 00:00:00 2001 From: TristonianJones Date: Thu, 10 Sep 2026 17:24:31 -0700 Subject: [PATCH 5/6] Updates to cost testing strategy --- common/cost/BUILD.bazel | 10 +- common/cost/aggregate_strategy.go | 16 +- common/cost/aggregate_strategy_test.go | 10 +- common/cost/estimator_test.go | 418 ++++--------------------- common/cost/tracker_test.go | 206 +++--------- 5 files changed, 126 insertions(+), 534 deletions(-) diff --git a/common/cost/BUILD.bazel b/common/cost/BUILD.bazel index b8e85e0e..cad08a76 100644 --- a/common/cost/BUILD.bazel +++ b/common/cost/BUILD.bazel @@ -8,6 +8,7 @@ package( go_library( name = "go_default_library", srcs = [ + "aggregate_strategy.go", "cost.go", "default_strategy.go", "estimator.go", @@ -31,6 +32,7 @@ go_test( name = "go_default_test", size = "small", srcs = [ + "aggregate_strategy_test.go", "cost_test.go", "default_strategy_test.go", "estimator_test.go", @@ -42,18 +44,14 @@ go_test( ":go_default_library", ], deps = [ - "//checker:go_default_library", - "//common:go_default_library", + "//cel:go_default_library", "//common/ast:go_default_library", - "//common/containers:go_default_library", "//common/decls:go_default_library", "//common/operators:go_default_library", "//common/overloads:go_default_library", - "//common/stdlib:go_default_library", "//common/types:go_default_library", "//common/types/ref:go_default_library", - "//interpreter:go_default_library", - "//parser:go_default_library", + "//ext:go_default_library", "//test/proto3pb:go_default_library", ], ) diff --git a/common/cost/aggregate_strategy.go b/common/cost/aggregate_strategy.go index 1874f277..5df2ec82 100644 --- a/common/cost/aggregate_strategy.go +++ b/common/cost/aggregate_strategy.go @@ -23,8 +23,20 @@ import ( // AggregateSizingStrategy returns a SizingStrategy that computes recursive size estimates // by following paths during cost estimation, and calculates actual runtime size using -// AggregateSize during cost tracking. By default, stringUnitLength is configured to 1 -// so that unscaled character/byte lengths are preserved for cost modeling. +// AggregateSize during cost tracking. +// +// The strategy pins stringUnitLength to 1, overriding the types package default, so that +// raw character and byte lengths reach the cost model unscaled. This is deliberate: sizing +// strategies report raw dimensions, and cost expressions own the conversion to cost units +// by applying factors such as StringTraversalCostFactor. +// +// Callers may supply additional SizeCalculatorOption values, which are applied after the +// pinned default. +// +// Warning: passing types.SizeCalculatorStringUnitLength with a value other than 1 overrides +// the pinned default and causes string costs to be discounted twice, once by the calculator +// and again by the cost factor in the cost expression. A unit length of 10 combined with +// StringTraversalCostFactor yields an effective 100x discount rather than 10x. func AggregateSizingStrategy(opts ...types.SizeCalculatorOption) SizingStrategy { if len(opts) == 0 { return defaultAggregateSizing diff --git a/common/cost/aggregate_strategy_test.go b/common/cost/aggregate_strategy_test.go index 29ca4cb7..afbf5cbe 100644 --- a/common/cost/aggregate_strategy_test.go +++ b/common/cost/aggregate_strategy_test.go @@ -172,11 +172,11 @@ func TestAggregateSizingStrategy_EstimateSize_List(t *testing.T) { } tests := []struct { - name string - ctx EstimateContext - node AstNode - wantOk bool - check func(t *testing.T, sz SizeEstimate) + name string + ctx EstimateContext + node AstNode + wantOk bool + check func(t *testing.T, sz SizeEstimate) }{ { name: "literal_list_elements", diff --git a/common/cost/estimator_test.go b/common/cost/estimator_test.go index 835ae5b3..cb226c77 100644 --- a/common/cost/estimator_test.go +++ b/common/cost/estimator_test.go @@ -20,23 +20,56 @@ import ( "strings" "testing" - "cel.dev/cel-go/checker" - "cel.dev/cel-go/common" + "cel.dev/cel-go/cel" "cel.dev/cel-go/common/ast" - "cel.dev/cel-go/common/containers" "cel.dev/cel-go/common/cost" "cel.dev/cel-go/common/decls" - "cel.dev/cel-go/common/operators" "cel.dev/cel-go/common/overloads" - "cel.dev/cel-go/common/stdlib" "cel.dev/cel-go/common/types" "cel.dev/cel-go/common/types/ref" - "cel.dev/cel-go/common/types/traits" - "cel.dev/cel-go/parser" + "cel.dev/cel-go/ext" proto3pb "cel.dev/cel-go/test/proto3pb" ) +// testCelEnv is the CEL environment shared by the cost estimator and tracker tests. +// +// Beyond the standard library it supplies the ext macros and functions these tests +// exercise: cel.bind from ext.Bindings, and the two-variable comprehensions (all, exists, +// existsOne, transformList, transformMap, transformMapEntry) from ext.TwoVarComprehensions +// along with the cel.@mapInsert function they expand to. These were previously hand-rolled +// in this file. +var testCelEnv = func() *cel.Env { + env, err := cel.NewEnv( + cel.Types(&proto3pb.TestAllTypes{}), + cel.CrossTypeNumericComparisons(true), + ext.Bindings(), + ext.TwoVarComprehensions(), + cel.Function("max", + cel.MemberOverload("list_bytes_max", + []*cel.Type{cel.ListType(cel.BytesType)}, cel.BytesType)), + ) + if err != nil { + panic(fmt.Sprintf("cel.NewEnv() failed: %v", err)) + } + return env +}() + +// compile type checks an expression against the shared test environment extended with the +// given variable declarations, and returns the checked AST. +func compile(t *testing.T, expr string, vars ...*decls.VariableDecl) *ast.AST { + t.Helper() + env, err := testCelEnv.Extend(cel.VariableDecls(vars...)) + if err != nil { + t.Fatalf("env.Extend() failed: %v", err) + } + checked, iss := env.Compile(expr) + if iss.Err() != nil { + t.Fatalf("env.Compile(%q) failed: %v", expr, iss.Err()) + } + return checked.NativeRep() +} + func TestCost(t *testing.T) { allTypes := types.NewObjectType("google.expr.proto3.test.TestAllTypes") allList := types.NewListType(allTypes) @@ -773,29 +806,29 @@ func TestCost(t *testing.T) { wanted: cost.CostEstimate{Min: 25, Max: 25}, }, { - name: "bind: with variable list and index", - expr: `cel.bind(a, input, a[0])`, - vars: []*decls.VariableDecl{decls.NewVariable("input", intList)}, + name: "bind: with variable list and index", + expr: `cel.bind(a, input, a[0])`, + vars: []*decls.VariableDecl{decls.NewVariable("input", intList)}, wanted: cost.CostEstimate{Min: 13, Max: 13}, }, { - name: "bind: with variable map and index", - expr: `cel.bind(m, input, m['key'])`, - vars: []*decls.VariableDecl{decls.NewVariable("input", types.NewMapType(types.StringType, types.StringType))}, + name: "bind: with variable map and index", + expr: `cel.bind(m, input, m['key'])`, + vars: []*decls.VariableDecl{decls.NewVariable("input", types.NewMapType(types.StringType, types.StringType))}, wanted: cost.CostEstimate{Min: 13, Max: 13}, }, { - name: "bind: with comprehension and size hints", - vars: []*decls.VariableDecl{decls.NewVariable("input", allList)}, - hints: map[string]uint64{"input": 100}, - expr: `cel.bind(a, input, a.all(x, true))`, + name: "bind: with comprehension and size hints", + vars: []*decls.VariableDecl{decls.NewVariable("input", allList)}, + hints: map[string]uint64{"input": 100}, + expr: `cel.bind(a, input, a.all(x, true))`, wanted: cost.CostEstimate{Min: 13, Max: 313}, }, { - name: "bind: nested with list and size hints", - vars: []*decls.VariableDecl{decls.NewVariable("input", nestedList)}, - hints: map[string]uint64{"input": 50, "input.@items": 10}, - expr: `cel.bind(a, input, a.all(x, x.all(y, true)))`, + name: "bind: nested with list and size hints", + vars: []*decls.VariableDecl{decls.NewVariable("input", nestedList)}, + hints: map[string]uint64{"input": 50, "input.@items": 10}, + expr: `cel.bind(a, input, a.all(x, x.all(y, true)))`, wanted: cost.CostEstimate{Min: 13, Max: 1763}, }, { @@ -912,44 +945,7 @@ func TestCost(t *testing.T) { if tc.hints == nil { tc.hints = map[string]uint64{} } - p, err := parser.NewParser(parser.Macros(testMacros...)) - if err != nil { - t.Fatalf("parser.NewParser() failed: %v", err) - } - src := common.NewStringSource(tc.expr, "") - pe, errs := p.Parse(src) - if len(errs.GetErrors()) != 0 { - t.Fatalf("parser.Parse(%v) failed: %v", tc.expr, errs.ToDisplayString()) - } - reg, err := types.NewRegistry(types.ProtoTypeDefs(&proto3pb.TestAllTypes{})) - if err != nil { - t.Fatalf("types.NewRegistry(...) failed: %v", err) - } - - e, err := checker.NewEnv(containers.DefaultContainer, reg) - if err != nil { - t.Fatalf("checker.NewEnv() failed: %v", err) - } - err = e.AddFunctions(stdlib.Functions()...) - if err != nil { - t.Fatalf("environment creation error: %v", err) - } - maxFunc, _ := decls.NewFunction("max", - decls.MemberOverload("list_bytes_max", - []*types.Type{types.NewListType(types.BytesType)}, - types.BytesType)) - err = e.AddFunctions(maxFunc, mapInsertFunctionDecl()) - if err != nil { - t.Fatalf("environment creation error: %v", err) - } - err = e.AddIdents(tc.vars...) - if err != nil { - t.Fatalf("environment creation error: %s\n", err) - } - checked, errs := checker.Check(pe, src, e) - if len(errs.GetErrors()) != 0 { - t.Fatalf("Check(%s) failed: %v", tc.expr, errs.ToDisplayString()) - } + checked := compile(t, tc.expr, tc.vars...) est, err := cost.Cost(checked, testCostEstimator{hints: tc.hints}, tc.options...) if err != nil { t.Fatalf("Cost() failed: %v", err) @@ -1040,35 +1036,8 @@ func (testCustomSizingStrategy) TrackSize(ctx cost.TrackContext, value ref.Val) } func TestCustomSizingStrategy(t *testing.T) { - prse, err := parser.NewParser(parser.Macros(parser.AllMacros...)) - if err != nil { - t.Fatalf("parser.NewParser() failed: %v", err) - } - src := common.NewStringSource("custom_str.contains('abc')", "") - pe, errs := prse.Parse(src) - if len(errs.GetErrors()) != 0 { - t.Fatalf("Parse() failed: %v", errs.ToDisplayString()) - } - reg, err := types.NewRegistry() - if err != nil { - t.Fatalf("types.NewRegistry() failed: %v", err) - } - e, err := checker.NewEnv(containers.DefaultContainer, reg) - if err != nil { - t.Fatalf("checker.NewEnv() failed: %v", err) - } - err = e.AddFunctions(stdlib.Functions()...) - if err != nil { - t.Fatalf("AddFunctions failed: %v", err) - } - err = e.AddIdents(decls.NewVariable("custom_str", types.StringType)) - if err != nil { - t.Fatalf("AddIdents failed: %v", err) - } - checked, errs := checker.Check(pe, src, e) - if len(errs.GetErrors()) != 0 { - t.Fatalf("Check() failed: %v", errs.ToDisplayString()) - } + checked := compile(t, "custom_str.contains('abc')", + decls.NewVariable("custom_str", types.StringType)) res, err := cost.Cost(checked, nil, cost.EstimateSizingStrategy(testCustomSizingStrategy{})) if err != nil { @@ -1083,276 +1052,3 @@ func TestCustomSizingStrategy(t *testing.T) { t.Errorf("got cost %v, wanted {Min: 2, Max: 3}", res) } } - -var ( - testMacros = append( - append([]parser.Macro{}, parser.AllMacros...), - testCelBindMacro(), - testTwoVarAllMacro(), - testTwoVarExistsMacro(), - testTwoVarExistsOneMacro(), - testTwoVarExistsOneMacroNew(), - testTwoVarTransformListMacro(), - testTwoVarTransformListFilterMacro(), - testTwoVarTransformMapMacro(), - testTwoVarTransformMapFilterMacro(), - testTwoVarTransformMapEntryMacro(), - testTwoVarTransformMapEntryFilterMacro(), - ) -) - -func testCelBindMacro() parser.Macro { - return parser.NewReceiverMacro("bind", 3, func(eh parser.ExprHelper, target ast.Expr, args []ast.Expr) (ast.Expr, *common.Error) { - if target == nil || target.Kind() != ast.IdentKind || target.AsIdent() != "cel" { - return nil, nil - } - varIdent := args[0] - if varIdent.Kind() != ast.IdentKind { - return nil, eh.NewError(varIdent.ID(), "cel.bind() variable names must be simple identifiers") - } - varName := varIdent.AsIdent() - varInit := args[1] - resultExpr := args[2] - return eh.NewComprehension( - eh.NewList(), - "#unused", - varName, - varInit, - eh.NewLiteral(types.False), - eh.NewIdent(varName), - resultExpr, - ), nil - }) -} - -func testTwoVarAllMacro() parser.Macro { - return parser.NewReceiverMacro("all", 3, func(eh parser.ExprHelper, target ast.Expr, args []ast.Expr) (ast.Expr, *common.Error) { - iterVar1, iterVar2, err := extractTwoVarIterVars(eh, args[0], args[1]) - if err != nil { - return nil, err - } - return eh.NewComprehensionTwoVar( - target, - iterVar1, - iterVar2, - eh.AccuIdentName(), - eh.NewLiteral(types.True), - eh.NewCall(operators.NotStrictlyFalse, eh.NewAccuIdent()), - eh.NewCall(operators.LogicalAnd, eh.NewAccuIdent(), args[2]), - eh.NewAccuIdent(), - ), nil - }) -} - -func testTwoVarExistsMacro() parser.Macro { - return parser.NewReceiverMacro("exists", 3, func(eh parser.ExprHelper, target ast.Expr, args []ast.Expr) (ast.Expr, *common.Error) { - iterVar1, iterVar2, err := extractTwoVarIterVars(eh, args[0], args[1]) - if err != nil { - return nil, err - } - return eh.NewComprehensionTwoVar( - target, - iterVar1, - iterVar2, - eh.AccuIdentName(), - eh.NewLiteral(types.False), - eh.NewCall(operators.NotStrictlyFalse, eh.NewCall(operators.LogicalNot, eh.NewAccuIdent())), - eh.NewCall(operators.LogicalOr, eh.NewAccuIdent(), args[2]), - eh.NewAccuIdent(), - ), nil - }) -} - -func testTwoVarExistsOneMacro() parser.Macro { - return parser.NewReceiverMacro("exists_one", 3, testTwoVarExistsOneExpander) -} - -func testTwoVarExistsOneMacroNew() parser.Macro { - return parser.NewReceiverMacro("existsOne", 3, testTwoVarExistsOneExpander) -} - -func testTwoVarExistsOneExpander(eh parser.ExprHelper, target ast.Expr, args []ast.Expr) (ast.Expr, *common.Error) { - iterVar1, iterVar2, err := extractTwoVarIterVars(eh, args[0], args[1]) - if err != nil { - return nil, err - } - return eh.NewComprehensionTwoVar( - target, - iterVar1, - iterVar2, - eh.AccuIdentName(), - eh.NewLiteral(types.Int(0)), - eh.NewLiteral(types.True), - eh.NewCall(operators.Conditional, args[2], - eh.NewCall(operators.Add, eh.NewAccuIdent(), eh.NewLiteral(types.Int(1))), - eh.NewAccuIdent()), - eh.NewCall(operators.Equals, eh.NewAccuIdent(), eh.NewLiteral(types.Int(1))), - ), nil -} - -func testTwoVarTransformListMacro() parser.Macro { - return parser.NewReceiverMacro("transformList", 3, testTwoVarTransformListExpander) -} - -func testTwoVarTransformListFilterMacro() parser.Macro { - return parser.NewReceiverMacro("transformList", 4, testTwoVarTransformListExpander) -} - -func testTwoVarTransformListExpander(eh parser.ExprHelper, target ast.Expr, args []ast.Expr) (ast.Expr, *common.Error) { - iterVar1, iterVar2, err := extractTwoVarIterVars(eh, args[0], args[1]) - if err != nil { - return nil, err - } - var transform, filter ast.Expr - if len(args) == 4 { - filter = args[2] - transform = args[3] - } else { - transform = args[2] - } - step := eh.NewCall(operators.Add, eh.NewAccuIdent(), eh.NewList(transform)) - if filter != nil { - step = eh.NewCall(operators.Conditional, filter, step, eh.NewAccuIdent()) - } - return eh.NewComprehensionTwoVar( - target, - iterVar1, - iterVar2, - eh.AccuIdentName(), - eh.NewList(), - eh.NewLiteral(types.True), - step, - eh.NewAccuIdent(), - ), nil -} - -func testTwoVarTransformMapMacro() parser.Macro { - return parser.NewReceiverMacro("transformMap", 3, testTwoVarTransformMapExpander) -} - -func testTwoVarTransformMapFilterMacro() parser.Macro { - return parser.NewReceiverMacro("transformMap", 4, testTwoVarTransformMapExpander) -} - -func testTwoVarTransformMapExpander(eh parser.ExprHelper, target ast.Expr, args []ast.Expr) (ast.Expr, *common.Error) { - iterVar1, iterVar2, err := extractTwoVarIterVars(eh, args[0], args[1]) - if err != nil { - return nil, err - } - var transform, filter ast.Expr - if len(args) == 4 { - filter = args[2] - transform = args[3] - } else { - transform = args[2] - } - step := eh.NewCall("cel.@mapInsert", eh.NewAccuIdent(), eh.NewIdent(iterVar1), transform) - if filter != nil { - step = eh.NewCall(operators.Conditional, filter, step, eh.NewAccuIdent()) - } - return eh.NewComprehensionTwoVar( - target, - iterVar1, - iterVar2, - eh.AccuIdentName(), - eh.NewMap(), - eh.NewLiteral(types.True), - step, - eh.NewAccuIdent(), - ), nil -} - -func testTwoVarTransformMapEntryMacro() parser.Macro { - return parser.NewReceiverMacro("transformMapEntry", 3, testTwoVarTransformMapEntryExpander) -} - -func testTwoVarTransformMapEntryFilterMacro() parser.Macro { - return parser.NewReceiverMacro("transformMapEntry", 4, testTwoVarTransformMapEntryExpander) -} - -func testTwoVarTransformMapEntryExpander(eh parser.ExprHelper, target ast.Expr, args []ast.Expr) (ast.Expr, *common.Error) { - iterVar1, iterVar2, err := extractTwoVarIterVars(eh, args[0], args[1]) - if err != nil { - return nil, err - } - var transform, filter ast.Expr - if len(args) == 4 { - filter = args[2] - transform = args[3] - } else { - transform = args[2] - } - step := eh.NewCall("cel.@mapInsert", eh.NewAccuIdent(), transform) - if filter != nil { - step = eh.NewCall(operators.Conditional, filter, step, eh.NewAccuIdent()) - } - return eh.NewComprehensionTwoVar( - target, - iterVar1, - iterVar2, - eh.AccuIdentName(), - eh.NewMap(), - eh.NewLiteral(types.True), - step, - eh.NewAccuIdent(), - ), nil -} - -func extractTwoVarIterVars(eh parser.ExprHelper, arg0, arg1 ast.Expr) (string, string, *common.Error) { - if arg0.Kind() != ast.IdentKind { - return "", "", eh.NewError(arg0.ID(), "argument must be a simple name") - } - if arg1.Kind() != ast.IdentKind { - return "", "", eh.NewError(arg1.ID(), "argument must be a simple name") - } - iterVar1 := arg0.AsIdent() - iterVar2 := arg1.AsIdent() - if iterVar1 == iterVar2 { - return "", "", eh.NewError(arg1.ID(), fmt.Sprintf("duplicate variable name: %s", iterVar1)) - } - if iterVar1 == eh.AccuIdentName() || iterVar1 == parser.AccumulatorName { - return "", "", eh.NewError(arg0.ID(), "iteration variable overwrites accumulator variable") - } - if iterVar2 == eh.AccuIdentName() || iterVar2 == parser.AccumulatorName { - return "", "", eh.NewError(arg1.ID(), "iteration variable overwrites accumulator variable") - } - return iterVar1, iterVar2, nil -} - -func mapInsertFunctionDecl() *decls.FunctionDecl { - kType := types.NewTypeParamType("K") - vType := types.NewTypeParamType("V") - mapKVType := types.NewMapType(kType, vType) - fn, _ := decls.NewFunction("cel.@mapInsert", - decls.Overload("@mapInsert_map_key_value", - []*types.Type{mapKVType, kType, vType}, - mapKVType), - decls.Overload("@mapInsert_map_map", - []*types.Type{mapKVType, mapKVType}, - mapKVType), - decls.SingletonFunctionBinding(func(args ...ref.Val) ref.Val { - if len(args) == 3 { - m := args[0].(traits.Mapper) - k := args[1] - v := args[2] - return types.InsertMapKeyValue(m, k, v) - } - if len(args) == 2 { - tm := args[0].(traits.Mapper) - um := args[1].(traits.Mapper) - umIt := um.Iterator() - for umIt.HasNext() == types.True { - k := umIt.Next() - updateOrErr := types.InsertMapKeyValue(tm, k, um.Get(k)) - if types.IsError(updateOrErr) { - return updateOrErr - } - tm = updateOrErr.(traits.Mapper) - } - return tm - } - return types.NoSuchOverloadErr() - }), - ) - return fn -} diff --git a/common/cost/tracker_test.go b/common/cost/tracker_test.go index dace09cb..459386d8 100644 --- a/common/cost/tracker_test.go +++ b/common/cost/tracker_test.go @@ -15,7 +15,6 @@ package cost_test import ( - "fmt" "math" "math/rand" "reflect" @@ -24,17 +23,12 @@ import ( "testing" "time" - "cel.dev/cel-go/checker" - "cel.dev/cel-go/common" - "cel.dev/cel-go/common/containers" + "cel.dev/cel-go/cel" "cel.dev/cel-go/common/cost" "cel.dev/cel-go/common/decls" "cel.dev/cel-go/common/overloads" - "cel.dev/cel-go/common/stdlib" "cel.dev/cel-go/common/types" "cel.dev/cel-go/common/types/ref" - "cel.dev/cel-go/interpreter" - "cel.dev/cel-go/parser" proto3pb "cel.dev/cel-go/test/proto3pb" ) @@ -277,14 +271,14 @@ func TestTrackCostAdvanced(t *testing.T) { ctx := constructActivation(t, tc.in) lhsCost, _, err := computeCost(t, tc.lhsExpr, nil, ctx, nil) if err != nil { - t.Fatalf("Interpreter.Eval(activation interpreter.Activation) failed to eval expression due: %v", err) + t.Fatalf("Program.Eval(activation) failed to eval expression due: %v", err) } rhsCost, _, err := computeCost(t, tc.rhsExpr, nil, ctx, nil) if err != nil { - t.Fatalf("Interpreter.Eval(activation interpreter.Activation) failed to eval expression due: %v", err) + t.Fatalf("Program.Eval(activation) failed to eval expression due: %v", err) } if lhsCost != rhsCost { - t.Errorf(`Interpreter.Eval(activation interpreter.Activation) failed return a cost for %s of %d equal to a cost for %s of %d`, + t.Errorf(`Program.Eval(activation) failed return a cost for %s of %d equal to a cost for %s of %d`, tc.lhsExpr, lhsCost, tc.rhsExpr, rhsCost) } }) @@ -313,93 +307,67 @@ func TestTrackCostAdvanced(t *testing.T) { ctx := constructActivation(t, tc.in) lhsCost, _, err := computeCost(t, tc.lhsExpr, nil, ctx, nil) if err != nil { - t.Fatalf("Interpreter.Eval(activation interpreter.Activation) failed to eval expression due: %v", err) + t.Fatalf("Program.Eval(activation) failed to eval expression due: %v", err) } rhsCost, _, err := computeCost(t, tc.rhsExpr, nil, ctx, nil) if err != nil { - t.Fatalf("Interpreter.Eval(activation interpreter.Activation) failed to eval expression due: %v", err) + t.Fatalf("Program.Eval(activation) failed to eval expression due: %v", err) } if lhsCost >= rhsCost { - t.Errorf(`Interpreter.Eval(activation interpreter.Activation) failed return a cost for %s of %d less than the cost for %s of %d`, + t.Errorf(`Program.Eval(activation) failed return a cost for %s of %d less than the cost for %s of %d`, tc.lhsExpr, lhsCost, tc.rhsExpr, rhsCost) } }) } } -func computeCost(t *testing.T, expr string, vars []*decls.VariableDecl, ctx interpreter.Activation, options []cost.TrackerOption) (actualCost uint64, est cost.CostEstimate, err error) { +func computeCost(t *testing.T, expr string, vars []*decls.VariableDecl, ctx cel.Activation, options []cost.TrackerOption) (actualCost uint64, est cost.CostEstimate, err error) { t.Helper() - s := common.NewTextSource(expr) - p, err := parser.NewParser(parser.Macros(testMacros...)) + env, err := testCelEnv.Extend(cel.VariableDecls(vars...)) if err != nil { - t.Fatalf("Failed to initialize parser: %v", err) + t.Fatalf("env.Extend() failed: %v", err) } - parsed, errs := p.Parse(s) - if len(errs.GetErrors()) != 0 { - t.Fatalf(`Failed to Parse expression "%s", error: %v`, expr, errs.GetErrors()) + checked, iss := env.Compile(expr) + if iss.Err() != nil { + t.Fatalf("env.Compile(%q) failed: %v", expr, iss.Err()) } - cont := containers.DefaultContainer - reg := newTestRegistry(t, types.ProtoTypeDefs(&proto3pb.TestAllTypes{})) - attrs := interpreter.NewAttributeFactory(cont, reg, reg) - env := newTestEnv(t, cont, reg) - err = env.AddFunctions(mapInsertFunctionDecl()) + // The estimate must be configured with the same presence test behavior as the tracker, + // so derive it from a tracker built with the same options. + tracker, err := cost.NewTracker(nil, options...) if err != nil { - t.Fatalf("Failed to add mapInsertFunctionDecl: %v", err) + t.Fatalf("cost.NewTracker() failed: %v", err) } - err = env.AddIdents(vars...) - if err != nil { - t.Fatalf("Failed to initialize env: %v", err) - } - costTracker, err := cost.NewTracker(&testRuntimeCostEstimator{}, options...) - if err != nil { - t.Fatalf("cost.NewCostTracker() failed: %v", err) - } - costTracker, err = costTracker.Clone() - if err != nil { - t.Fatalf("checker.Clone() 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, err = cost.Cost(checked, testTrackerCostEstimator{}, cost.PresenceTestHasCost(costTracker.PresenceTestHasCost())) + est, err = cost.Cost(checked.NativeRep(), testTrackerCostEstimator{}, + cost.PresenceTestHasCost(tracker.PresenceTestHasCost())) if err != nil { t.Fatalf("cost.Cost() failed: %v", err) } - interp := newStandardInterpreter(t, cont, reg, reg, attrs, mapInsertFunctionDecl()) - prg, err := interp.NewInterpretable(checked, - interpreter.CostObserver(interpreter.CostTrackerFactory(func() (*cost.Tracker, error) { - return costTracker, nil - }))) + + prg, err := env.Program(checked, + cel.CostTracking(&testRuntimeCostEstimator{}), + cel.CostTrackerOptions(options...)) if err != nil { - t.Fatalf(`Failed to check expression "%s", error: %v`, expr, errs.GetErrors()) + t.Fatalf("env.Program() failed: %v", err) } - - defer func() { - if r := recover(); r != nil { - switch t := r.(type) { - case interpreter.EvalCancelledError: - err = t - default: - err = fmt.Errorf("internal error: %v", r) - } - } - }() - frame := interpreter.AsFrame(ctx) - prg.Exec(frame) - return costTracker.ActualCost(), est, err + // Program.Eval recovers evaluation panics itself, and attaches the cost tracker to the + // details even when evaluation fails, so a cost limit breach still reports its cost. + _, det, err := prg.Eval(ctx) + if cost := det.ActualCost(); cost != nil { + actualCost = *cost + } + return actualCost, est, err } -func constructActivation(t testing.TB, in any) interpreter.Activation { +func constructActivation(t testing.TB, in any) cel.Activation { t.Helper() if in == nil { - return interpreter.EmptyActivation() + return cel.NoVars() } - a, err := interpreter.NewActivation(in) + a, err := cel.NewActivation(in) if err != nil { - t.Fatalf("interpreter.NewActivation(%v) failed: %v", in, err) + t.Fatalf("cel.NewActivation(%v) failed: %v", in, err) } return a } @@ -1268,16 +1236,16 @@ func TestRuntimeCost(t *testing.T) { if tc.expectExceedsLimit { return } - t.Fatalf("Interpreter.Eval(activation interpreter.Activation) failed due to: %v", err) + t.Fatalf("Program.Eval(activation) failed due to: %v", err) } if tc.expectExceedsLimit { - t.Fatalf("Interpreter.Eval(activation interpreter.Activation) failed to return a cost exceeded error for limit %d, got cost %d", tc.limit, actualCost) + t.Fatalf("Program.Eval(activation) failed to return a cost exceeded error for limit %d, got cost %d", tc.limit, actualCost) } if actualCost != tc.want { - t.Fatalf("Interpreter.Eval(activation interpreter.Activation) failed to return expected runtime cost %d, got %d", tc.want, actualCost) + t.Fatalf("Program.Eval(activation) failed to return expected runtime cost %d, got %d", tc.want, actualCost) } if est.Min > actualCost || est.Max < actualCost { - t.Fatalf("Interpreter.Eval(activation interpreter.Activation) failed to return cost in range of estimate cost [%d, %d], got %d", + t.Fatalf("Program.Eval(activation) failed to return cost in range of estimate cost [%d, %d], got %d", est.Min, est.Max, actualCost) } }) @@ -1315,42 +1283,17 @@ func BenchmarkCostTracking(b *testing.B) { for _, bm := range benchmarks { b.Run(bm.name, func(b *testing.B) { - s := common.NewTextSource(bm.expr) - p, err := parser.NewParser(parser.Macros(parser.AllMacros...)) - if err != nil { - b.Fatalf("Failed to initialize parser: %v", err) - } - parsed, errs := p.Parse(s) - if len(errs.GetErrors()) != 0 { - b.Fatalf("Parse(%s) failed: %v", bm.expr, errs.GetErrors()) - } - - cont := containers.DefaultContainer - reg := newTestRegistry(b, types.ProtoTypeDefs(&proto3pb.TestAllTypes{})) - attrs := interpreter.NewAttributeFactory(cont, reg, reg) - env := newTestEnv(b, cont, reg) - if len(bm.vars) > 0 { - err = env.AddIdents(bm.vars...) - if err != nil { - b.Fatalf("Failed to add idents: %v", err) - } - } - checked, errs := checker.Check(parsed, s, env) - if len(errs.GetErrors()) != 0 { - b.Fatalf("Check(%s) failed: %v", bm.expr, errs.GetErrors()) - } - - evalCostTracker, err := cost.NewTracker(nil) + env, err := testCelEnv.Extend(cel.VariableDecls(bm.vars...)) if err != nil { - b.Fatalf("cost.NewCostTracker() failed: %v", err) + b.Fatalf("env.Extend() failed: %v", err) } - trackerFactory := func() (*cost.Tracker, error) { - return evalCostTracker.Clone() + checked, iss := env.Compile(bm.expr) + if iss.Err() != nil { + b.Fatalf("env.Compile(%q) failed: %v", bm.expr, iss.Err()) } - interp := newStandardInterpreter(b, cont, reg, reg, attrs) - prg, err := interp.NewInterpretable(checked, interpreter.CostObserver(interpreter.CostTrackerFactory(trackerFactory))) + prg, err := env.Program(checked, cel.CostTracking(nil)) if err != nil { - b.Fatalf("NewInterpretable(%s) failed: %v", bm.expr, err) + b.Fatalf("env.Program(%s) failed: %v", bm.expr, err) } ctx := constructActivation(b, bm.in) @@ -1363,63 +1306,6 @@ func BenchmarkCostTracking(b *testing.B) { } } -func newTestEnv(t testing.TB, cont *containers.Container, reg *types.Registry) *checker.Env { - t.Helper() - env, err := checker.NewEnv(cont, reg, checker.CrossTypeNumericComparisons(true)) - if err != nil { - t.Fatalf("checker.NewEnv(%v, %v) failed: %v", cont, reg, err) - } - err = env.AddFunctions(stdlib.Functions()...) - if err != nil { - t.Fatalf("env.Add(stdlib.Functions()...) failed: %v", err) - } - return env -} - -func newTestRegistry(t testing.TB, opts ...types.RegistryOption) *types.Registry { - t.Helper() - var o []any - for _, opt := range opts { - o = append(o, opt) - } - reg, err := types.NewRegistry(o...) - if err != nil { - t.Fatalf("types.NewRegistry() failed: %v", err) - } - return reg -} - -func newStandardInterpreter(t testing.TB, - container *containers.Container, - provider types.Provider, - adapter types.Adapter, - resolver interpreter.AttributeFactory, - optFuncs ...*decls.FunctionDecl) interpreter.Interpreter { - t.Helper() - disp := interpreter.NewDispatcher() - for _, fn := range stdlib.Functions() { - bindings, err := fn.Bindings() - if err != nil { - t.Fatalf("fn.Bindings() failed for function %v. error: %v", fn.Name(), err) - } - err = disp.Add(bindings...) - if err != nil { - t.Fatalf("dispatcher.Add() failed: %v", err) - } - } - for _, fn := range optFuncs { - bindings, err := fn.Bindings() - if err != nil { - t.Fatalf("fn.Bindings() failed for function %v. error: %v", fn.Name(), err) - } - err = disp.Add(bindings...) - if err != nil { - t.Fatalf("dispatcher.Add() failed: %v", err) - } - } - return interpreter.NewInterpreter(disp, container, provider, adapter, resolver) -} - type testConcurrentSizingStrategy struct{} func (testConcurrentSizingStrategy) EstimateSize(ctx cost.EstimateContext, node cost.AstNode) (cost.SizeEstimate, bool) { From 6751b947ecb498a860004950dfec7a863b5a2116 Mon Sep 17 00:00:00 2001 From: TristonianJones Date: Thu, 10 Sep 2026 17:55:20 -0700 Subject: [PATCH 6/6] Validated the that updates in estimation match the actual cost already present in v0.28.1+ --- common/cost/estimator.go | 18 +++++++++++---- common/cost/estimator_test.go | 36 +++++++++++++++++++++++++++++ common/cost/tracker_test.go | 43 +++++++++++++++++++++++++++++++++++ 3 files changed, 93 insertions(+), 4 deletions(-) diff --git a/common/cost/estimator.go b/common/cost/estimator.go index 1e974c7d..d01b3771 100644 --- a/common/cost/estimator.go +++ b/common/cost/estimator.go @@ -362,14 +362,19 @@ func (c *coster) relativeAttributeCost(operand ast.Expr) CostEstimate { // isAttributeChain reports whether an expression is resolved as part of a single attribute // during evaluation. A chain begins at an identifier, or at a ternary which selects between -// attributes, and is extended by field selections and index operations. +// attributes, and is extended by field selections and index operations, including their +// optional variants. +// +// This predicate must mirror the planner's decision to wrap an operand in a relative +// attribute: any expression which plans to an interpretable attribute extends the chain and +// does not incur an additional attribute resolution cost. func isAttributeChain(e ast.Expr) bool { switch e.Kind() { case ast.IdentKind, ast.SelectKind: return true case ast.CallKind: switch e.AsCall().FunctionName() { - case operators.Index, operators.Conditional: + case operators.Index, operators.OptIndex, operators.OptSelect, operators.Conditional: return true } } @@ -388,8 +393,13 @@ func (c *coster) costCall(e ast.Expr) CostEstimate { args := call.Args() var sum CostEstimate - if call.FunctionName() == operators.Index && len(args) > 0 { - sum = sum.Add(c.relativeAttributeCost(args[0])) + // Index-like operators qualify their first argument, which requires a relative attribute + // when the operand is a computed value rather than a named one. + switch call.FunctionName() { + case operators.Index, operators.OptIndex, operators.OptSelect: + if len(args) > 0 { + sum = sum.Add(c.relativeAttributeCost(args[0])) + } } argTypes := make([]AstNode, len(args)) diff --git a/common/cost/estimator_test.go b/common/cost/estimator_test.go index cb226c77..ab6bd5a5 100644 --- a/common/cost/estimator_test.go +++ b/common/cost/estimator_test.go @@ -43,6 +43,7 @@ var testCelEnv = func() *cel.Env { env, err := cel.NewEnv( cel.Types(&proto3pb.TestAllTypes{}), cel.CrossTypeNumericComparisons(true), + cel.OptionalTypes(), ext.Bindings(), ext.TwoVarComprehensions(), cel.Function("max", @@ -677,6 +678,41 @@ func TestCost(t *testing.T) { expr: `['hello', 'hi'][0] != ['hello', 'bye'][1]`, wanted: cost.CostEstimate{Min: 25, Max: 25}, }, + { + // Optional index over a computed operand costs the same as its non-optional + // counterpart: the planner qualifies a relative attribute in both cases. + name: "literal map optional access", + expr: `{'hello': 'hi'}[?'hello']`, + wanted: cost.CostEstimate{Min: 32, Max: 32}, + }, + { + name: "literal map optional select", + expr: `{'hello': 'hi'}.?hello`, + wanted: cost.CostEstimate{Min: 32, Max: 32}, + }, + { + name: "literal list optional access", + expr: `['hello', 'hi'][?0]`, + wanted: cost.CostEstimate{Min: 12, Max: 12}, + }, + { + // An optional select extends the attribute chain, so the trailing selection + // must not be charged an extra attribute resolution. + name: "optional select chain", + expr: `self.?val1.val2`, + vars: []*decls.VariableDecl{ + decls.NewVariable("self", types.NewMapType(types.StringType, types.DynType)), + }, + wanted: cost.CostEstimate{Min: 3, Max: 3}, + }, + { + name: "optional index chain", + expr: `self[?'val1'].val2`, + vars: []*decls.VariableDecl{ + decls.NewVariable("self", types.NewMapType(types.StringType, types.DynType)), + }, + wanted: cost.CostEstimate{Min: 3, Max: 3}, + }, { name: "type call", expr: `type(1)`, diff --git a/common/cost/tracker_test.go b/common/cost/tracker_test.go index 459386d8..c60b0065 100644 --- a/common/cost/tracker_test.go +++ b/common/cost/tracker_test.go @@ -436,6 +436,7 @@ func TestRuntimeCost(t *testing.T) { allMap := types.NewMapType(types.StringType, allTypes) nestedMap := types.NewMapType(types.StringType, allMap) + nestedMapStr := types.NewMapType(types.StringType, types.NewMapType(types.StringType, types.StringType)) cases := []struct { name string expr string @@ -503,6 +504,48 @@ func TestRuntimeCost(t *testing.T) { want: 3, in: map[string]any{"input": []string{"v"}}, }, + { + name: "optional select: map", + expr: `input.?key`, + vars: []*decls.VariableDecl{decls.NewVariable("input", types.NewMapType(types.StringType, types.StringType))}, + want: 2, + in: map[string]any{"input": map[string]string{"key": "v"}}, + }, + { + name: "optional index: map", + expr: `input[?'key']`, + vars: []*decls.VariableDecl{decls.NewVariable("input", types.NewMapType(types.StringType, types.StringType))}, + want: 2, + in: map[string]any{"input": map[string]string{"key": "v"}}, + }, + { + // An optional select extends the attribute chain, so the trailing selection + // only adds a qualifier cost rather than a second attribute resolution. + name: "optional select: chained", + expr: `input.?key.subkey`, + vars: []*decls.VariableDecl{decls.NewVariable("input", nestedMapStr)}, + want: 3, + in: map[string]any{"input": map[string]map[string]string{"key": {"subkey": "v"}}}, + }, + { + name: "optional index: chained", + expr: `input[?'key'].subkey`, + vars: []*decls.VariableDecl{decls.NewVariable("input", nestedMapStr)}, + want: 3, + in: map[string]any{"input": map[string]map[string]string{"key": {"subkey": "v"}}}, + }, + { + // A computed operand requires a relative attribute, which costs an extra + // attribute resolution on top of the qualifier. + name: "optional index: map literal", + expr: `{'key': 'v'}[?'key']`, + want: 32, + }, + { + name: "optional index: list literal", + expr: `['v'][?0]`, + want: 12, + }, { name: "select: field test only no has() cost", expr: `has(input.single_int32)`,