| // Copyright 2023 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 cel_test |
| |
| import ( |
| "sort" |
| "strings" |
| "testing" |
| |
| "github.com/google/cel-go/cel" |
| "github.com/google/cel-go/common/ast" |
| "github.com/google/cel-go/ext" |
| |
| "google.golang.org/protobuf/encoding/prototext" |
| "google.golang.org/protobuf/proto" |
| |
| proto3pb "github.com/google/cel-go/test/proto3pb" |
| |
| exprpb "google.golang.org/genproto/googleapis/api/expr/v1alpha1" |
| ) |
| |
| func TestStaticOptimizerUpdateExpr(t *testing.T) { |
| expr := `has(a.b)` |
| inlined := `[x, y].filter(i, i.size() > 0)[0].z` |
| |
| e := optimizerEnv(t) |
| exprAST, iss := e.Compile(expr) |
| if iss.Err() != nil { |
| t.Fatalf("Compile() failed: %v", iss.Err()) |
| } |
| |
| inlinedAST, iss := e.Compile(inlined) |
| if iss.Err() != nil { |
| t.Fatalf("Compile() failed: %v", iss.Err()) |
| } |
| opt, err := cel.NewStaticOptimizer(&testOptimizer{t: t, inlineExpr: inlinedAST.NativeRep()}) |
| if err != nil { |
| t.Fatalf("NewStaticOptimizer() failed: %v", err) |
| } |
| optAST, iss := opt.Optimize(e, exprAST) |
| if iss.Err() != nil { |
| t.Fatalf("Optimize() generated an invalid AST: %v", iss.Err()) |
| } |
| optString, err := cel.AstToString(optAST) |
| if err != nil { |
| t.Fatalf("cel.AstToString() failed: %v", err) |
| } |
| sourceInfo := optAST.NativeRep().SourceInfo() |
| sourceInfoPB, err := ast.SourceInfoToProto(sourceInfo) |
| if err != nil { |
| t.Fatalf("cel.AstToCheckedExpr() failed: %v", err) |
| } |
| sourceInfoPB.Positions = nil |
| wantTextPB := ` |
| location: "<input>" |
| line_offsets: 9 |
| macro_calls: { |
| key: 1 |
| value: { |
| call_expr: { |
| function: "has" |
| args: { |
| id: 24 |
| select_expr: { |
| operand: { |
| id: 2 |
| call_expr: { |
| function: "_[_]" |
| args: { |
| id: 3 |
| } |
| args: { |
| id: 20 |
| const_expr: { |
| int64_value: 0 |
| } |
| } |
| } |
| } |
| field: "z" |
| } |
| } |
| } |
| } |
| } |
| macro_calls: { |
| key: 3 |
| value: { |
| call_expr: { |
| target: { |
| id: 4 |
| list_expr: { |
| elements: { |
| id: 5 |
| ident_expr: { |
| name: "x" |
| } |
| } |
| elements: { |
| id: 6 |
| ident_expr: { |
| name: "y" |
| } |
| } |
| } |
| } |
| function: "filter" |
| args: { |
| id: 17 |
| ident_expr: { |
| name: "i" |
| } |
| } |
| args: { |
| id: 10 |
| call_expr: { |
| function: "_>_" |
| args: { |
| id: 11 |
| call_expr: { |
| target: { |
| id: 12 |
| ident_expr: { |
| name: "i" |
| } |
| } |
| function: "size" |
| } |
| } |
| args: { |
| id: 13 |
| const_expr: { |
| int64_value: 0 |
| } |
| } |
| } |
| } |
| } |
| } |
| } |
| ` |
| var wantSourceInfoPB exprpb.SourceInfo |
| if err := prototext.Unmarshal([]byte(wantTextPB), &wantSourceInfoPB); err != nil { |
| t.Fatalf("prototext.Unmarshal() failed: %v", err) |
| } |
| if !proto.Equal(&wantSourceInfoPB, sourceInfoPB) { |
| t.Errorf("got source info: %s, wanted %s", prototext.Format(sourceInfoPB), wantTextPB) |
| } |
| expected := `has([x, y].filter(i, i.size() > 0)[0].z)` |
| if expected != optString { |
| t.Errorf("inlined got %q, wanted %q", optString, expected) |
| } |
| } |
| |
| func TestStaticOptimizerNewAST(t *testing.T) { |
| tests := []string{ |
| `[3, 2, 1]`, |
| `[1, 2, 3].all(i, i != 0)`, |
| `cel.bind(m, {"a": 1, "b": 2}, m.filter(k, m[k] > 1))`, |
| } |
| for _, tst := range tests { |
| tc := tst |
| t.Run(tc, func(t *testing.T) { |
| e := optimizerEnv(t) |
| exprAST, iss := e.Compile(tc) |
| if iss.Err() != nil { |
| t.Fatalf("Compile(%q) failed: %v", tc, iss.Err()) |
| } |
| opt, err := cel.NewStaticOptimizer(&identityOptimizer{t: t}) |
| if err != nil { |
| t.Fatalf("NewStaticOptimizer() failed: %v", err) |
| } |
| optAST, iss := opt.Optimize(e, exprAST) |
| if iss.Err() != nil { |
| t.Fatalf("Optimize() generated an invalid AST: %v", iss.Err()) |
| } |
| optString, err := cel.AstToString(optAST) |
| if err != nil { |
| t.Fatalf("cel.AstToString() failed: %v", err) |
| } |
| if tc != optString { |
| t.Errorf("identity optimizer got %q, wanted %q", optString, tc) |
| } |
| }) |
| } |
| } |
| |
| func TestOptimizeWithSource(t *testing.T) { |
| initial := `has(a.b)` |
| replacement := `x["a"]` |
| e := optimizerEnv(t) |
| initialAST, iss := e.Compile(initial) |
| if iss.Err() != nil { |
| t.Fatalf("Compile(%q) failed: %v", initial, iss.Err()) |
| } |
| replacementAST, iss := e.Compile(replacement) |
| if iss.Err() != nil { |
| t.Fatalf("Compile(%q) failed: %v", replacement, iss.Err()) |
| } |
| |
| opt, err := cel.NewStaticOptimizer( |
| &replaceOptimizer{t: t, targetAST: replacementAST.NativeRep()}, |
| cel.OptimizeWithSource(replacementAST.Source()), |
| ) |
| if err != nil { |
| t.Fatalf("NewStaticOptimizer() failed: %v", err) |
| } |
| optAST, iss := opt.Optimize(e, initialAST) |
| if iss.Err() != nil { |
| t.Fatalf("Optimize() returned an error: %v", iss.Err()) |
| } |
| |
| if optAST.Source().Content() != replacement { |
| t.Errorf("got source content %q, wanted %q", optAST.Source().Content(), replacement) |
| } |
| sourceInfoPB, err := ast.SourceInfoToProto(optAST.NativeRep().SourceInfo()) |
| if err != nil { |
| t.Fatalf("cel.AstToCheckedExpr() failed: %v", err) |
| } |
| wantTextPB := ` |
| location: "<input>" |
| line_offsets: 7 |
| positions: { |
| key: 1 |
| value: 1 |
| } |
| positions: { |
| key: 2 |
| value: 0 |
| } |
| positions: { |
| key: 3 |
| value: 2 |
| } |
| ` |
| var wantSourceInfoPB exprpb.SourceInfo |
| if err := prototext.Unmarshal([]byte(wantTextPB), &wantSourceInfoPB); err != nil { |
| t.Fatalf("prototext.Unmarshal() failed: %v", err) |
| } |
| if !proto.Equal(&wantSourceInfoPB, sourceInfoPB) { |
| t.Errorf("got source info: %s, wanted %s", prototext.Format(sourceInfoPB), wantTextPB) |
| } |
| } |
| |
| func TestStaticOptimizerNilAST(t *testing.T) { |
| env := optimizerEnv(t) |
| opt, err := cel.NewStaticOptimizer(&identityOptimizer{t: t}) |
| if err != nil { |
| t.Fatalf("NewStaticOptimizer() failed: %v", err) |
| } |
| optAST, iss := opt.Optimize(env, nil) |
| if iss.Err() == nil || !strings.Contains(iss.Err().Error(), "unexpected unspecified type") { |
| t.Errorf("opt.Optimize(env, nil) got (%v, %v), wanted unexpected unspecified type", optAST, iss) |
| } |
| } |
| |
| type identityOptimizer struct { |
| t *testing.T |
| } |
| |
| func (opt *identityOptimizer) Optimize(ctx *cel.OptimizerContext, a *ast.AST) *ast.AST { |
| opt.t.Helper() |
| // The copy method should effectively update all of the old macro refs with new ones that are |
| // identical, but renumbered. |
| main := ctx.CopyASTAndMetadata(a) |
| // The new AST call will create a parsed expression which will be type-checked by the static |
| // optimizer. The input and output expressions should be identical, though may vary by number |
| // though. |
| return ctx.NewAST(main) |
| } |
| |
| type testOptimizer struct { |
| t *testing.T |
| inlineExpr *ast.AST |
| } |
| |
| func (opt *testOptimizer) Optimize(ctx *cel.OptimizerContext, a *ast.AST) *ast.AST { |
| opt.t.Helper() |
| copy := ctx.CopyASTAndMetadata(opt.inlineExpr) |
| origID := a.Expr().ID() |
| presenceTest, hasMacro := ctx.NewHasMacro(origID, copy) |
| macroKeys := getMacroKeys(ctx.MacroCalls()) |
| if len(macroKeys) != 2 { |
| opt.t.Errorf("Got %v macro calls, wanted 2", macroKeys) |
| } |
| ctx.UpdateExpr(a.Expr(), presenceTest) |
| ctx.SetMacroCall(origID, hasMacro) |
| return ctx.NewAST(a.Expr()) |
| } |
| |
| type replaceOptimizer struct { |
| t *testing.T |
| targetAST *ast.AST |
| } |
| |
| func (opt *replaceOptimizer) Optimize(ctx *cel.OptimizerContext, a *ast.AST) *ast.AST { |
| opt.t.Helper() |
| copy := ctx.CopyASTAndMetadata(opt.targetAST) |
| return ctx.NewAST(copy) |
| } |
| |
| func getMacroKeys(macroCalls map[int64]ast.Expr) []int { |
| keys := []int{} |
| for k := range macroCalls { |
| keys = append(keys, int(k)) |
| } |
| sort.Ints(keys) |
| return keys |
| } |
| |
| func optimizerEnv(t *testing.T) *cel.Env { |
| t.Helper() |
| opts := []cel.EnvOption{ |
| cel.Types(&proto3pb.TestAllTypes{}), |
| cel.OptionalTypes(), |
| cel.EnableMacroCallTracking(), |
| ext.Bindings(), |
| cel.Variable("a", cel.MapType(cel.StringType, cel.StringType)), |
| cel.Variable("x", cel.MapType(cel.StringType, cel.StringType)), |
| cel.Variable("y", cel.MapType(cel.StringType, cel.StringType)), |
| } |
| e, err := cel.NewEnv(opts...) |
| if err != nil { |
| t.Fatalf("NewEnv() failed: %v", err) |
| } |
| return e |
| } |