Add cost tracking for encoders

ext/README.md

@@ -33,6 +33,8 @@ Encoding utilities for marshalling data into standardized representations.
 
 ### Base64.Decode
 
+**Introduced in version 0 (cost support in version 1)**
+
 Decodes base64-encoded string to bytes.
 
 This function will return an error if the string input is not
@@ -47,6 +49,8 @@ Examples:
 
 ### Base64.Encode
 
+**Introduced in version 0 (cost support in version 1)**
+
 Encodes bytes to a base64-encoded string.
 
     base64.encode(<bytes>)  -> <string>

ext/encoders.go

@@ -19,8 +19,10 @@ import (
 	"math"
 
 	"github.com/google/cel-go/cel"
+	"github.com/google/cel-go/checker"
 	"github.com/google/cel-go/common/types"
 	"github.com/google/cel-go/common/types/ref"
+	"github.com/google/cel-go/interpreter"
 )
 
 // Encoders returns a cel.EnvOption to configure extended functions for string, byte, and object
@@ -75,8 +77,8 @@ func (*encoderLib) LibraryName() string {
 	return "cel.lib.ext.encoders"
 }
 
-func (*encoderLib) CompileOptions() []cel.EnvOption {
-	return []cel.EnvOption{
+func (lib *encoderLib) CompileOptions() []cel.EnvOption {
+	opts := []cel.EnvOption{
 		cel.Function("base64.decode",
 			cel.Overload("base64_decode_string", []*cel.Type{cel.StringType}, cel.BytesType,
 				cel.UnaryBinding(func(str ref.Val) ref.Val {
@@ -90,10 +92,26 @@ func (*encoderLib) CompileOptions() []cel.EnvOption {
 					return stringOrError(base64EncodeBytes([]byte(b)))
 				}))),
 	}
+	if lib.version >= 1 {
+		estimators := []checker.CostOption{
+			checker.OverloadCostEstimate("base64_decode_string", estimateDecode),
+			checker.OverloadCostEstimate("base64_encode_bytes", estimateEncode),
+		}
+		opts = append(opts, cel.CostEstimatorOptions(estimators...))
+	}
+	return opts
 }
 
-func (*encoderLib) ProgramOptions() []cel.ProgramOption {
-	return []cel.ProgramOption{}
+func (lib *encoderLib) ProgramOptions() []cel.ProgramOption {
+	var opts []cel.ProgramOption
+	if lib.version >= 1 {
+		trackers := []interpreter.CostTrackerOption{
+			interpreter.OverloadCostTracker("base64_decode_string", trackDecode),
+			interpreter.OverloadCostTracker("base64_encode_bytes", trackEncode),
+		}
+		opts = append(opts, cel.CostTrackerOptions(trackers...))
+	}
+	return opts
 }
 
 func base64DecodeString(str string) ([]byte, error) {
@@ -110,3 +128,50 @@ func base64DecodeString(str string) ([]byte, error) {
 func base64EncodeBytes(bytes []byte) (string, error) {
 	return base64.StdEncoding.EncodeToString(bytes), nil
 }
+
+func estimateEncode(estimator checker.CostEstimator, target *checker.AstNode, args []checker.AstNode) *checker.CallEstimate {
+	if len(args) != 1 {
+		return nil
+	}
+	sz := estimateSize(estimator, args[0])
+	cost := sz.MultiplyByCostFactor(stringCostFactor).Add(callCostEstimate)
+	resSize := estimateEncodeSize(sz)
+	return &checker.CallEstimate{CostEstimate: cost, ResultSize: &resSize}
+}
+
+func estimateDecode(estimator checker.CostEstimator, target *checker.AstNode, args []checker.AstNode) *checker.CallEstimate {
+	if len(args) != 1 {
+		return nil
+	}
+	sz := estimateSize(estimator, args[0])
+	cost := sz.MultiplyByCostFactor(stringCostFactor).Add(callCostEstimate)
+	resSize := estimateDecodeSize(sz)
+	return &checker.CallEstimate{CostEstimate: cost, ResultSize: &resSize}
+}
+
+func trackEncode(args []ref.Val, _ ref.Val) *uint64 {
+	sz := actualSize(args[0])
+	cost := uint64(math.Ceil(float64(sz)*stringCostFactor)) + callCost
+	return &cost
+}
+
+func trackDecode(args []ref.Val, _ ref.Val) *uint64 {
+	sz := actualSize(args[0])
+	cost := uint64(math.Ceil(float64(sz)*stringCostFactor)) + callCost
+	return &cost
+}
+
+func estimateEncodeSize(sz checker.SizeEstimate) checker.SizeEstimate {
+	minVal := (sz.Min*4 + 2) / 3
+	maxVal := (sz.Max*4 + 2) / 3
+	if sz.Max > math.MaxUint64/4 {
+		maxVal = math.MaxUint64
+	}
+	return checker.SizeEstimate{Min: minVal, Max: maxVal}
+}
+
+func estimateDecodeSize(sz checker.SizeEstimate) checker.SizeEstimate {
+	minVal := sz.Min * 3 / 4
+	maxVal := sz.Max * 3 / 4
+	return checker.SizeEstimate{Min: minVal, Max: maxVal}
+}

ext/encoders_test.go

@@ -20,6 +20,7 @@ import (
 	"testing"
 
 	"github.com/google/cel-go/cel"
+	"github.com/google/cel-go/checker"
 )
 
 func TestEncoders(t *testing.T) {
@@ -93,3 +94,128 @@ func TestEncodersVersion(t *testing.T) {
 		t.Fatalf("EncodersVersion(0) failed: %v", err)
 	}
 }
+
+func testEncodersCostsEnv(t *testing.T, version int, opts ...cel.EnvOption) *cel.Env {
+	t.Helper()
+	baseOpts := []cel.EnvOption{
+		Encoders(EncodersVersion(uint32(version))),
+		cel.EnableMacroCallTracking(),
+	}
+	env, err := cel.NewEnv(append(baseOpts, opts...)...)
+	if err != nil {
+		t.Fatalf("cel.NewEnv(Encoders()) failed: %v", err)
+	}
+	return env
+}
+
+func TestEncodersCosts(t *testing.T) {
+	tests := []struct {
+		name          string
+		expr          string
+		vars          []cel.EnvOption
+		in            map[string]any
+		hints         map[string]uint64
+		estimatedCost checker.CostEstimate
+		actualCost    uint64
+		version       int
+	}{
+		{
+			name: "encode_bytes_v0",
+			expr: "base64.encode(x) == 'aGVsbG8='",
+			vars: []cel.EnvOption{
+				cel.Variable("x", cel.BytesType),
+			},
+			in: map[string]any{
+				"x": []byte("hello"),
+			},
+			hints: map[string]uint64{
+				"x": 100,
+			},
+			estimatedCost: checker.FixedCostEstimate(3), // x lookup (1) + encode (1) + == (1) = 3
+			actualCost:    3,
+			version:       0,
+		},
+		{
+			name: "encode_bytes_v1",
+			expr: "base64.encode(x) == 'aGVsbG8='",
+			vars: []cel.EnvOption{
+				cel.Variable("x", cel.BytesType),
+			},
+			in: map[string]any{
+				"x": []byte("hello"),
+			},
+			hints: map[string]uint64{
+				"x": 100,
+			},
+			estimatedCost: checker.CostEstimate{Min: 3, Max: 13}, // x lookup (1) + encode (100 * 0.1 + 1 = 11) + == (1) = 13
+			actualCost:    4,                                     // x lookup (1) + encode (ceil(5 * 0.1) + 1 = 2) + == (1) = 4
+			version:       1,
+		},
+		{
+			name: "decode_string_v0",
+			expr: "base64.decode(x) == b'hello'",
+			vars: []cel.EnvOption{
+				cel.Variable("x", cel.StringType),
+			},
+			in: map[string]any{
+				"x": "aGVsbG8=",
+			},
+			hints: map[string]uint64{
+				"x": 100,
+			},
+			estimatedCost: checker.FixedCostEstimate(3),
+			actualCost:    3,
+			version:       0,
+		},
+		{
+			name: "decode_string_v1",
+			expr: "base64.decode(x) == b'hello'",
+			vars: []cel.EnvOption{
+				cel.Variable("x", cel.StringType),
+			},
+			in: map[string]any{
+				"x": "aGVsbG8=",
+			},
+			hints: map[string]uint64{
+				"x": 100,
+			},
+			estimatedCost: checker.CostEstimate{Min: 3, Max: 13}, // x lookup (1) + decode (100 * 0.1 + 1 = 11) + == (1) = 13
+			actualCost:    4,                                     // x lookup (1) + decode (ceil(8 * 0.1) + 1 = 2) + == (1) = 4
+			version:       1,
+		},
+		{
+			name:          "encode_bytes_v1_literal",
+			expr:          "base64.encode(b'hello') == 'aGVsbG8='",
+			estimatedCost: checker.FixedCostEstimate(3),
+			actualCost:    3,
+			version:       1,
+		},
+		{
+			name:          "decode_string_v1_literal",
+			expr:          "base64.decode('aGVsbG8=') == b'hello'",
+			estimatedCost: checker.FixedCostEstimate(3),
+			actualCost:    3,
+			version:       1,
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			env := testEncodersCostsEnv(t, tc.version, tc.vars...)
+			var asts []*cel.Ast
+			pAst, iss := env.Parse(tc.expr)
+			if iss.Err() != nil {
+				t.Fatalf("env.Parse(%v) failed: %v", tc.expr, iss.Err())
+			}
+			asts = append(asts, pAst)
+			cAst, iss := env.Check(pAst)
+			if iss.Err() != nil {
+				t.Fatalf("env.Check(%v) failed: %v", tc.expr, iss.Err())
+			}
+			testCheckCost(t, env, cAst, tc.hints, tc.estimatedCost)
+			asts = append(asts, cAst)
+			for _, ast := range asts {
+				testEvalWithCost(t, env, ast, tc.in, tc.actualCost)
+			}
+		})
+	}
+}