Skip to content

Commit 7b55cc0

Browse files
TristonianJonesTristonianJones
committed
Add cost estimation for encoders
1 parent f0ffa7e commit 7b55cc0

3 files changed

Lines changed: 199 additions & 4 deletions

File tree

‎ext/README.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,8 @@ Encoding utilities for marshalling data into standardized representations.
3333

3434
### Base64.Decode
3535

36+
**Introduced in version 0 (cost support in version 1)**
37+
3638
Decodes base64-encoded string to bytes.
3739

3840
This function will return an error if the string input is not
@@ -47,6 +49,8 @@ Examples:
4749

4850
### Base64.Encode
4951

52+
**Introduced in version 0 (cost support in version 1)**
53+
5054
Encodes bytes to a base64-encoded string.
5155

5256
base64.encode(<bytes>) -> <string>

‎ext/encoders.go‎

Lines changed: 69 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,10 @@ import (
1919
"math"
2020

2121
"github.com/google/cel-go/cel"
22+
"github.com/google/cel-go/checker"
2223
"github.com/google/cel-go/common/types"
2324
"github.com/google/cel-go/common/types/ref"
25+
"github.com/google/cel-go/interpreter"
2426
)
2527

2628
// Encoders returns a cel.EnvOption to configure extended functions for string, byte, and object
@@ -75,8 +77,8 @@ func (*encoderLib) LibraryName() string {
7577
return "cel.lib.ext.encoders"
7678
}
7779

78-
func (*encoderLib) CompileOptions() []cel.EnvOption {
79-
return []cel.EnvOption{
80+
func (lib *encoderLib) CompileOptions() []cel.EnvOption {
81+
opts := []cel.EnvOption{
8082
cel.Function("base64.decode",
8183
cel.Overload("base64_decode_string", []*cel.Type{cel.StringType}, cel.BytesType,
8284
cel.UnaryBinding(func(str ref.Val) ref.Val {
@@ -90,10 +92,26 @@ func (*encoderLib) CompileOptions() []cel.EnvOption {
9092
return stringOrError(base64EncodeBytes([]byte(b)))
9193
}))),
9294
}
95+
if lib.version >= 1 {
96+
estimators := []checker.CostOption{
97+
checker.OverloadCostEstimate("base64_decode_string", estimateDecode),
98+
checker.OverloadCostEstimate("base64_encode_bytes", estimateEncode),
99+
}
100+
opts = append(opts, cel.CostEstimatorOptions(estimators...))
101+
}
102+
return opts
93103
}
94104

95-
func (*encoderLib) ProgramOptions() []cel.ProgramOption {
96-
return []cel.ProgramOption{}
105+
func (lib *encoderLib) ProgramOptions() []cel.ProgramOption {
106+
var opts []cel.ProgramOption
107+
if lib.version >= 1 {
108+
trackers := []interpreter.CostTrackerOption{
109+
interpreter.OverloadCostTracker("base64_decode_string", trackDecode),
110+
interpreter.OverloadCostTracker("base64_encode_bytes", trackEncode),
111+
}
112+
opts = append(opts, cel.CostTrackerOptions(trackers...))
113+
}
114+
return opts
97115
}
98116

99117
func base64DecodeString(str string) ([]byte, error) {
@@ -110,3 +128,50 @@ func base64DecodeString(str string) ([]byte, error) {
110128
func base64EncodeBytes(bytes []byte) (string, error) {
111129
return base64.StdEncoding.EncodeToString(bytes), nil
112130
}
131+
132+
func estimateEncode(estimator checker.CostEstimator, target *checker.AstNode, args []checker.AstNode) *checker.CallEstimate {
133+
if len(args) != 1 {
134+
return nil
135+
}
136+
sz := estimateSize(estimator, args[0])
137+
cost := sz.MultiplyByCostFactor(stringCostFactor).Add(callCostEstimate)
138+
resSize := estimateEncodeSize(sz)
139+
return &checker.CallEstimate{CostEstimate: cost, ResultSize: &resSize}
140+
}
141+
142+
func estimateDecode(estimator checker.CostEstimator, target *checker.AstNode, args []checker.AstNode) *checker.CallEstimate {
143+
if len(args) != 1 {
144+
return nil
145+
}
146+
sz := estimateSize(estimator, args[0])
147+
cost := sz.MultiplyByCostFactor(stringCostFactor).Add(callCostEstimate)
148+
resSize := estimateDecodeSize(sz)
149+
return &checker.CallEstimate{CostEstimate: cost, ResultSize: &resSize}
150+
}
151+
152+
func trackEncode(args []ref.Val, _ ref.Val) *uint64 {
153+
sz := actualSize(args[0])
154+
cost := uint64(math.Ceil(float64(sz)*stringCostFactor)) + callCost
155+
return &cost
156+
}
157+
158+
func trackDecode(args []ref.Val, _ ref.Val) *uint64 {
159+
sz := actualSize(args[0])
160+
cost := uint64(math.Ceil(float64(sz)*stringCostFactor)) + callCost
161+
return &cost
162+
}
163+
164+
func estimateEncodeSize(sz checker.SizeEstimate) checker.SizeEstimate {
165+
minVal := (sz.Min*4 + 2) / 3
166+
maxVal := (sz.Max*4 + 2) / 3
167+
if sz.Max > math.MaxUint64/4 {
168+
maxVal = math.MaxUint64
169+
}
170+
return checker.SizeEstimate{Min: minVal, Max: maxVal}
171+
}
172+
173+
func estimateDecodeSize(sz checker.SizeEstimate) checker.SizeEstimate {
174+
minVal := sz.Min * 3 / 4
175+
maxVal := sz.Max * 3 / 4
176+
return checker.SizeEstimate{Min: minVal, Max: maxVal}
177+
}

‎ext/encoders_test.go‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020
"testing"
2121

2222
"github.com/google/cel-go/cel"
23+
"github.com/google/cel-go/checker"
2324
)
2425

2526
func TestEncoders(t *testing.T) {
@@ -93,3 +94,128 @@ func TestEncodersVersion(t *testing.T) {
9394
t.Fatalf("EncodersVersion(0) failed: %v", err)
9495
}
9596
}
97+
98+
func testEncodersCostsEnv(t *testing.T, version int, opts ...cel.EnvOption) *cel.Env {
99+
t.Helper()
100+
baseOpts := []cel.EnvOption{
101+
Encoders(EncodersVersion(uint32(version))),
102+
cel.EnableMacroCallTracking(),
103+
}
104+
env, err := cel.NewEnv(append(baseOpts, opts...)...)
105+
if err != nil {
106+
t.Fatalf("cel.NewEnv(Encoders()) failed: %v", err)
107+
}
108+
return env
109+
}
110+
111+
func TestEncodersCosts(t *testing.T) {
112+
tests := []struct {
113+
name string
114+
expr string
115+
vars []cel.EnvOption
116+
in map[string]any
117+
hints map[string]uint64
118+
estimatedCost checker.CostEstimate
119+
actualCost uint64
120+
version int
121+
}{
122+
{
123+
name: "encode_bytes_v0",
124+
expr: "base64.encode(x) == 'aGVsbG8='",
125+
vars: []cel.EnvOption{
126+
cel.Variable("x", cel.BytesType),
127+
},
128+
in: map[string]any{
129+
"x": []byte("hello"),
130+
},
131+
hints: map[string]uint64{
132+
"x": 100,
133+
},
134+
estimatedCost: checker.FixedCostEstimate(3), // x lookup (1) + encode (1) + == (1) = 3
135+
actualCost: 3,
136+
version: 0,
137+
},
138+
{
139+
name: "encode_bytes_v1",
140+
expr: "base64.encode(x) == 'aGVsbG8='",
141+
vars: []cel.EnvOption{
142+
cel.Variable("x", cel.BytesType),
143+
},
144+
in: map[string]any{
145+
"x": []byte("hello"),
146+
},
147+
hints: map[string]uint64{
148+
"x": 100,
149+
},
150+
estimatedCost: checker.CostEstimate{Min: 3, Max: 13}, // x lookup (1) + encode (100 * 0.1 + 1 = 11) + == (1) = 13
151+
actualCost: 4, // x lookup (1) + encode (ceil(5 * 0.1) + 1 = 2) + == (1) = 4
152+
version: 1,
153+
},
154+
{
155+
name: "decode_string_v0",
156+
expr: "base64.decode(x) == b'hello'",
157+
vars: []cel.EnvOption{
158+
cel.Variable("x", cel.StringType),
159+
},
160+
in: map[string]any{
161+
"x": "aGVsbG8=",
162+
},
163+
hints: map[string]uint64{
164+
"x": 100,
165+
},
166+
estimatedCost: checker.FixedCostEstimate(3),
167+
actualCost: 3,
168+
version: 0,
169+
},
170+
{
171+
name: "decode_string_v1",
172+
expr: "base64.decode(x) == b'hello'",
173+
vars: []cel.EnvOption{
174+
cel.Variable("x", cel.StringType),
175+
},
176+
in: map[string]any{
177+
"x": "aGVsbG8=",
178+
},
179+
hints: map[string]uint64{
180+
"x": 100,
181+
},
182+
estimatedCost: checker.CostEstimate{Min: 3, Max: 13}, // x lookup (1) + decode (100 * 0.1 + 1 = 11) + == (1) = 13
183+
actualCost: 4, // x lookup (1) + decode (ceil(8 * 0.1) + 1 = 2) + == (1) = 4
184+
version: 1,
185+
},
186+
{
187+
name: "encode_bytes_v1_literal",
188+
expr: "base64.encode(b'hello') == 'aGVsbG8='",
189+
estimatedCost: checker.FixedCostEstimate(3),
190+
actualCost: 3,
191+
version: 1,
192+
},
193+
{
194+
name: "decode_string_v1_literal",
195+
expr: "base64.decode('aGVsbG8=') == b'hello'",
196+
estimatedCost: checker.FixedCostEstimate(3),
197+
actualCost: 3,
198+
version: 1,
199+
},
200+
}
201+
for _, tc := range tests {
202+
t.Run(tc.name, func(t *testing.T) {
203+
env := testEncodersCostsEnv(t, tc.version, tc.vars...)
204+
var asts []*cel.Ast
205+
pAst, iss := env.Parse(tc.expr)
206+
if iss.Err() != nil {
207+
t.Fatalf("env.Parse(%v) failed: %v", tc.expr, iss.Err())
208+
}
209+
asts = append(asts, pAst)
210+
cAst, iss := env.Check(pAst)
211+
if iss.Err() != nil {
212+
t.Fatalf("env.Check(%v) failed: %v", tc.expr, iss.Err())
213+
}
214+
testCheckCost(t, env, cAst, tc.hints, tc.estimatedCost)
215+
asts = append(asts, cAst)
216+
for _, ast := range asts {
217+
testEvalWithCost(t, env, ast, tc.in, tc.actualCost)
218+
}
219+
})
220+
}
221+
}

0 commit comments

Comments
 (0)