From 5a0dc476be99cf705770c9bde84f370cee6a8fd0 Mon Sep 17 00:00:00 2001 From: Anders Eknert Date: Thu, 18 Dec 2025 12:04:18 +0100 Subject: [PATCH] Template string performance improvements and more (#8143) A mixed bag of improvements I have had around for a while, and would like to see included in v1.12.0 if possible. I have about twice the amount of changes **not** included here as they could use some more testing. These changes should be low risk, I believe.. but obviously do let me know if you see any potential risks that I don't! - Add template string benchmarks - Faster template strings / print eval by not passing bctx in recursion - Allow passing nil value to `Query.WithQueryTracer` (no-op) - Reduce allocations in rego v1 compiler stages - Intern a few more common var name `Value`s - Remove redundant switch on `scope` in annotations code - Add a few more benchmarks in the `ast` package - Performance improvements in type checker, most notably removing a function literal for checking expression, which only ever had one implementation. We can extend this later if needed. - Prefer `NewGenericTransformer` over `&GenericTransformer` for easier tracking in pprof Signed-off-by: Anders Eknert --- v1/ast/annotations.go | 13 +--- v1/ast/check.go | 120 ++++++++++++----------------- v1/ast/env.go | 65 ++++++---------- v1/ast/interning.go | 15 +++- v1/ast/policy.go | 2 +- v1/ast/rego_v1.go | 19 +++-- v1/ast/strings.go | 2 + v1/ast/term_bench_test.go | 47 +++++++++++ v1/ast/transform.go | 14 ++-- v1/bundle/store.go | 2 +- v1/format/format.go | 4 +- v1/repl/repl.go | 3 +- v1/topdown/print.go | 22 +++--- v1/topdown/query.go | 2 +- v1/topdown/template_string.go | 5 +- v1/topdown/template_string_test.go | 110 +++++++++++++++----------- v1/topdown/topdown_bench_test.go | 63 +++++++++++++++ 17 files changed, 297 insertions(+), 211 deletions(-) diff --git a/v1/ast/annotations.go b/v1/ast/annotations.go index 47c543004c..603ab5cd77 100644 --- a/v1/ast/annotations.go +++ b/v1/ast/annotations.go @@ -433,18 +433,7 @@ func (a *Annotations) toObject() (*Object, *Error) { } if len(a.Scope) > 0 { - switch a.Scope { - case annotationScopeDocument: - obj.Insert(InternedTerm("scope"), InternedTerm("document")) - case annotationScopePackage: - obj.Insert(InternedTerm("scope"), InternedTerm("package")) - case annotationScopeRule: - obj.Insert(InternedTerm("scope"), InternedTerm("rule")) - case annotationScopeSubpackages: - obj.Insert(InternedTerm("scope"), InternedTerm("subpackages")) - default: - obj.Insert(InternedTerm("scope"), StringTerm(a.Scope)) - } + obj.Insert(InternedTerm("scope"), InternedTerm(a.Scope)) } if len(a.Title) > 0 { diff --git a/v1/ast/check.go b/v1/ast/check.go index ef04c39f00..e67599b4d2 100644 --- a/v1/ast/check.go +++ b/v1/ast/check.go @@ -7,7 +7,6 @@ package ast import ( "fmt" "slices" - "sort" "strings" "github.com/open-policy-agent/opa/v1/types" @@ -16,11 +15,6 @@ import ( type varRewriter func(Ref) Ref -// exprChecker defines the interface for executing type checking on a single -// expression. The exprChecker must update the provided TypeEnv with inferred -// types of vars. -type exprChecker func(*TypeEnv, *Expr) *Error - // typeChecker implements type checking on queries and rules. Errors are // accumulated on the typeChecker so that a single run can report multiple // issues. @@ -28,7 +22,6 @@ type typeChecker struct { builtins map[string]*Builtin required *Capabilities errs Errors - exprCheckers map[string]exprChecker varRewriter varRewriter ss *SchemaSet allowNet []string @@ -39,11 +32,7 @@ type typeChecker struct { // newTypeChecker returns a new typeChecker object that has no errors. func newTypeChecker() *typeChecker { - return &typeChecker{ - exprCheckers: map[string]exprChecker{ - "eq": checkExprEq, - }, - } + return &typeChecker{} } func (tc *typeChecker) newEnv(exist *TypeEnv) *TypeEnv { @@ -132,35 +121,31 @@ func (tc *typeChecker) CheckBody(env *TypeEnv, body Body) (*TypeEnv, Errors) { vis := newRefChecker(env, tc.varRewriter) gv := NewGenericVisitor(vis.Visit) - WalkExprs(body, func(expr *Expr) bool { - closureErrs := tc.checkClosures(env, expr) - for _, err := range closureErrs { - errors = append(errors, err) - } + for _, bexpr := range body { + WalkExprs(bexpr, func(expr *Expr) bool { + closureErrs := tc.checkClosures(env, expr) + errors = append(errors, closureErrs...) - hasClosureErrors := len(closureErrs) > 0 + // reset errors from previous iteration + vis.errs = nil + gv.Walk(expr) + errors = append(errors, vis.errs...) - // reset errors from previous iteration - vis.errs = nil - gv.Walk(expr) - for _, err := range vis.errs { - errors = append(errors, err) - } - - hasRefErrors := len(vis.errs) > 0 - - if err := tc.checkExpr(env, expr); err != nil { - // Suppress this error if a more actionable one has occurred. In - // this case, if an error occurred in a ref or closure contained in - // this expression, and the error is due to a nil type, then it's - // likely to be the result of the more specific error. - skip := (hasClosureErrors || hasRefErrors) && causedByNilType(err) - if !skip { - errors = append(errors, err) + if err := tc.checkExpr(env, expr); err != nil { + hasClosureErrors := len(closureErrs) > 0 + hasRefErrors := len(vis.errs) > 0 + // Suppress this error if a more actionable one has occurred. In + // this case, if an error occurred in a ref or closure contained in + // this expression, and the error is due to a nil type, then it's + // likely to be the result of the more specific error. + skip := (hasClosureErrors || hasRefErrors) && causedByNilType(err) + if !skip { + errors = append(errors, err) + } } - } - return true - }) + return true + }) + } tc.err(errors...) return env, errors @@ -281,23 +266,25 @@ func (tc *typeChecker) checkRule(env *TypeEnv, as *AnnotationSet, rule *Rule) { var tpe types.Type if len(rule.Head.Args) > 0 { - // If args are not referred to in body, infer as any. - WalkVars(rule.Head.Args, func(v Var) bool { - if cpy.GetByValue(v) == nil { - cpy.tree.PutOne(v, types.A) - } - return false - }) + for _, arg := range rule.Head.Args { + // If args are not referred to in body, infer as any. + WalkTerms(arg, func(t *Term) bool { + if _, ok := t.Value.(Var); ok { + if cpy.GetByValue(t.Value) == nil { + cpy.tree.PutOne(t.Value, types.A) + } + } + return false + }) + } // Construct function type. args := make([]types.Type, len(rule.Head.Args)) - for i := range len(rule.Head.Args) { + for i := range rule.Head.Args { args[i] = cpy.GetByValue(rule.Head.Args[i].Value) } - f := types.NewFunction(args, cpy.Get(rule.Head.Value)) - - tpe = f + tpe = types.NewFunction(args, cpy.GetByValue(rule.Head.Value.Value)) } else { switch rule.Head.RuleKind() { case SingleValue: @@ -374,9 +361,8 @@ func (tc *typeChecker) checkExpr(env *TypeEnv, expr *Expr) *Error { } } - checker := tc.exprCheckers[operator] - if checker != nil { - return checker(env, expr) + if operator == "eq" { + return checkExprEq(env, expr) } return tc.checkExprBuiltin(env, expr) @@ -599,7 +585,7 @@ func unify1(env *TypeEnv, term *Term, tpe types.Type, union bool) bool { return unifies } return false - case Set: + case *set: switch tpe := tpe.(type) { case *types.Set: return unify1Set(env, v, tpe, union) @@ -674,7 +660,7 @@ func unify1Object(env *TypeEnv, val Object, tpe *types.Object, union bool) bool return !stop } -func unify1Set(env *TypeEnv, val Set, tpe *types.Set, union bool) bool { +func unify1Set(env *TypeEnv, val *set, tpe *types.Set, union bool) bool { of := types.Values(tpe) return !val.Until(func(elem *Term) bool { return !unify1(env, elem, of, union) @@ -702,7 +688,6 @@ func newRefChecker(env *TypeEnv, f varRewriter) *refChecker { return &refChecker{ env: env, - errs: nil, varRewriter: f, } } @@ -806,7 +791,6 @@ func (rc *refChecker) checkRef(curr *TypeEnv, node *typeTreeNode, ref Ref, idx i } func (rc *refChecker) checkRefLeaf(tpe types.Type, ref Ref, idx int) *Error { - if idx == len(ref) { return nil } @@ -821,16 +805,16 @@ func (rc *refChecker) checkRefLeaf(tpe types.Type, ref Ref, idx int) *Error { switch value := head.Value.(type) { case Var: - if exist := rc.env.GetByValue(value); exist != nil { + if exist := rc.env.GetByValue(head.Value); exist != nil { if !unifies(exist, keys) { return newRefErrInvalid(ref[0].Location, rc.varRewriter(ref), idx, exist, keys, getOneOfForType(tpe)) } } else { - rc.env.tree.PutOne(value, types.Keys(tpe)) + rc.env.tree.PutOne(head.Value, types.Keys(tpe)) } case Ref: - if exist := rc.env.Get(value); exist != nil { + if exist := rc.env.GetByRef(value); exist != nil { if !unifies(exist, keys) { return newRefErrInvalid(ref[0].Location, rc.varRewriter(ref), idx, exist, keys, getOneOfForType(tpe)) } @@ -1131,7 +1115,7 @@ func getOneOfForNode(node *typeTreeNode) (result []Value) { return false }) - sortValueSlice(result) + slices.SortFunc(result, Value.Compare) return result } @@ -1154,16 +1138,10 @@ func getOneOfForType(tpe types.Type) (result []Value) { } result = removeDuplicate(result) - sortValueSlice(result) + slices.SortFunc(result, Value.Compare) return result } -func sortValueSlice(sl []Value) { - sort.Slice(sl, func(i, j int) bool { - return sl[i].Compare(sl[j]) < 0 - }) -} - func removeDuplicate(list []Value) []Value { seen := make(map[Value]bool) var newResult []Value @@ -1187,13 +1165,13 @@ func getArgTypes(env *TypeEnv, args []*Term) []types.Type { // getPrefix returns the shortest prefix of ref that exists in env func getPrefix(env *TypeEnv, ref Ref) (Ref, types.Type) { if len(ref) == 1 { - t := env.Get(ref) + t := env.GetByRef(ref) if t != nil { return ref, t } } for i := 1; i < len(ref); i++ { - t := env.Get(ref[:i]) + t := env.GetByRef(ref[:i]) if t != nil { return ref[:i], t } @@ -1201,12 +1179,14 @@ func getPrefix(env *TypeEnv, ref Ref) (Ref, types.Type) { return nil, nil } +var dynamicAnyAny = types.NewDynamicProperty(types.A, types.A) + // override takes a type t and returns a type obtained from t where the path represented by ref within it has type o (overriding the original type of that path) func override(ref Ref, t types.Type, o types.Type, rule *Rule) (types.Type, *Error) { var newStaticProps []*types.StaticProperty obj, ok := t.(*types.Object) if !ok { - newType, err := getObjectType(ref, o, rule, types.NewDynamicProperty(types.A, types.A)) + newType, err := getObjectType(ref, o, rule, dynamicAnyAny) if err != nil { return nil, err } diff --git a/v1/ast/env.go b/v1/ast/env.go index cb01d5299b..b55f32df0a 100644 --- a/v1/ast/env.go +++ b/v1/ast/env.go @@ -61,8 +61,7 @@ func (env *TypeEnv) GetByValue(v Value) types.Type { case *Array: static := make([]types.Type, x.Len()) for i := range static { - tpe := env.GetByValue(x.Elem(i).Value) - static[i] = tpe + static[i] = env.GetByValue(x.Elem(i).Value) } var dynamic types.Type @@ -80,17 +79,13 @@ func (env *TypeEnv) GetByValue(v Value) types.Type { x.Foreach(func(k, v *Term) { if IsConstant(k.Value) { - kjson, err := JSON(k.Value) - if err == nil { - tpe := env.GetByValue(v.Value) - static = append(static, types.NewStaticProperty(kjson, tpe)) + if kjson, err := JSON(k.Value); err == nil { + static = append(static, types.NewStaticProperty(kjson, env.GetByValue(v.Value))) return } } // Can't handle it as a static property, fallback to dynamic - typeK := env.GetByValue(k.Value) - typeV := env.GetByValue(v.Value) - dynamic = types.NewDynamicProperty(typeK, typeV) + dynamic = types.NewDynamicProperty(env.GetByValue(k.Value), env.GetByValue(v.Value)) }) if len(static) == 0 && dynamic == nil { @@ -99,7 +94,7 @@ func (env *TypeEnv) GetByValue(v Value) types.Type { return types.NewObject(static, dynamic) - case Set: + case *set: var tpe types.Type x.Foreach(func(elem *Term) { tpe = types.Or(tpe, env.GetByValue(elem.Value)) @@ -162,7 +157,6 @@ func (env *TypeEnv) GetByRef(ref Ref) types.Type { } func (env *TypeEnv) getRefFallback(ref Ref) types.Type { - if env.next != nil { return env.next.GetByRef(ref) } @@ -299,15 +293,11 @@ func (n *typeTreeNode) PutOne(key Value, tpe types.Type) { func (n *typeTreeNode) Put(path Ref, tpe types.Type) { curr := n for _, term := range path { - c, ok := curr.children.Get(term.Value) - - var child *typeTreeNode + child, ok := curr.children.Get(term.Value) if !ok { child = newTypeTree() child.key = term.Value curr.children.Put(child.key, child) - } else { - child = c } curr = child @@ -321,23 +311,18 @@ func (n *typeTreeNode) Put(path Ref, tpe types.Type) { func (n *typeTreeNode) Insert(path Ref, tpe types.Type, env *TypeEnv) { curr := n for i, term := range path { - c, ok := curr.children.Get(term.Value) - - var child *typeTreeNode + child, ok := curr.children.Get(term.Value) if !ok { child = newTypeTree() child.key = term.Value curr.children.Put(child.key, child) - } else { - child = c - if child.value != nil && i+1 < len(path) { - // If child has an object value, merge the new value into it. - if o, ok := child.value.(*types.Object); ok { - var err error - child.value, err = insertIntoObject(o, path[i+1:], tpe, env) - if err != nil { - panic(fmt.Errorf("unreachable, insertIntoObject: %w", err)) - } + } else if child.value != nil && i+1 < len(path) { + // If child has an object value, merge the new value into it. + if o, ok := child.value.(*types.Object); ok { + var err error + child.value, err = insertIntoObject(o, path[i+1:], tpe, env) + if err != nil { + panic(fmt.Errorf("unreachable, insertIntoObject: %w", err)) } } } @@ -349,8 +334,7 @@ func (n *typeTreeNode) Insert(path Ref, tpe types.Type, env *TypeEnv) { if _, ok := tpe.(*types.Object); ok && curr.children.Len() > 0 { // merge all leafs into the inserted object - leafs := curr.Leafs() - for p, t := range leafs { + for p, t := range curr.Leafs() { var err error curr.value, err = insertIntoObject(curr.value.(*types.Object), *p, t, env) if err != nil { @@ -388,7 +372,8 @@ func mergeTypes(a, b types.Type) types.Type { bDynProps := bObj.DynamicProperties() dynProps := types.NewDynamicProperty( types.Or(aDynProps.Key, bDynProps.Key), - mergeTypes(aDynProps.Value, bDynProps.Value)) + mergeTypes(aDynProps.Value, bDynProps.Value), + ) return types.NewObject(nil, dynProps) } else if bAny, ok := b.(types.Any); ok && len(a.StaticProperties()) == 0 { // If a is an object type with no static components ... @@ -417,14 +402,14 @@ func mergeTypes(a, b types.Type) types.Type { } func (n *typeTreeNode) String() string { - b := strings.Builder{} + b := &strings.Builder{} + key := "-" if k := n.key; k != nil { - b.WriteString(k.String()) - } else { - b.WriteString("-") + key = k.String() } + b.WriteString(key) if v := n.value; v != nil { b.WriteString(": ") b.WriteString(v.String()) @@ -432,9 +417,7 @@ func (n *typeTreeNode) String() string { n.children.Iter(func(_ Value, child *typeTreeNode) bool { b.WriteString("\n\t+ ") - s := child.String() - s = strings.ReplaceAll(s, "\n", "\n\t") - b.WriteString(s) + b.WriteString(strings.ReplaceAll(child.String(), "\n", "\n\t")) return false }) @@ -485,7 +468,8 @@ func (n *typeTreeNode) Leafs() map[*Ref]types.Type { func collectLeafs(n *typeTreeNode, path Ref, leafs map[*Ref]types.Type) { nPath := append(path, NewTerm(n.key)) if n.Leaf() { - leafs[&nPath] = n.Value() + npc := nPath // copy of else nPath escapes to heap even if !n.Leaf() + leafs[&npc] = n.Value() return } n.children.Iter(func(_ Value, v *typeTreeNode) bool { @@ -513,7 +497,6 @@ func selectConstant(tpe types.Type, term *Term) types.Type { // contains vars or refs, then the returned type will be a union of the // possible types. func selectRef(tpe types.Type, ref Ref) types.Type { - if tpe == nil || len(ref) == 0 { return tpe } diff --git a/v1/ast/interning.go b/v1/ast/interning.go index 6c722c6eff..67dd0ed49b 100644 --- a/v1/ast/interning.go +++ b/v1/ast/interning.go @@ -42,10 +42,17 @@ var ( } internedVarValues = map[string]Value{ - "input": Var("input"), - "data": Var("data"), - "key": Var("key"), - "value": Var("value"), + "input": Var("input"), + "data": Var("data"), + "args": Var("args"), + "schema": Var("schema"), + "key": Var("key"), + "value": Var("value"), + "future": Var("future"), + "rego": Var("rego"), + "set": Var("set"), + "internal": Var("internal"), + "else": Var("else"), "i": Var("i"), "j": Var("j"), "k": Var("k"), "v": Var("v"), "x": Var("x"), "y": Var("y"), "z": Var("z"), } diff --git a/v1/ast/policy.go b/v1/ast/policy.go index 512c1f454b..b87a1be4a1 100644 --- a/v1/ast/policy.go +++ b/v1/ast/policy.go @@ -621,7 +621,7 @@ func (imp *Import) SetLoc(loc *Location) { // document. This is the alias if defined otherwise the last element in the // path. func (imp *Import) Name() Var { - if len(imp.Alias) != 0 { + if imp.Alias != "" { return imp.Alias } switch v := imp.Path.Value.(type) { diff --git a/v1/ast/rego_v1.go b/v1/ast/rego_v1.go index 2880ea5d38..db9e0f722c 100644 --- a/v1/ast/rego_v1.go +++ b/v1/ast/rego_v1.go @@ -27,13 +27,12 @@ func checkRootDocumentOverrides(node any) Errors { errors := Errors{} WalkRules(node, func(rule *Rule) bool { - var name string + name := rule.Head.Name if len(rule.Head.Reference) > 0 { - name = rule.Head.Reference[0].Value.(Var).String() - } else { - name = rule.Head.Name.String() + name = rule.Head.Reference[0].Value.(Var) } - if RootDocumentRefs.Contains(RefTerm(VarTerm(name))) { + + if ReservedVars.Contains(name) { errors = append(errors, NewError(CompileErr, rule.Location, "rules must not shadow %v (use a different rule name)", name)) } @@ -52,8 +51,8 @@ func checkRootDocumentOverrides(node any) Errors { if expr.IsAssignment() { // assign() can be called directly, so we need to assert its given first operand exists before checking its name. if nameOp := expr.Operand(0); nameOp != nil { - name := nameOp.String() - if RootDocumentRefs.Contains(RefTerm(VarTerm(name))) { + name := Var(nameOp.String()) + if ReservedVars.Contains(name) { errors = append(errors, NewError(CompileErr, expr.Location, "variables must not shadow %v (use a different variable name)", name)) } } @@ -66,16 +65,16 @@ func checkRootDocumentOverrides(node any) Errors { func walkCalls(node any, f func(any) bool) { vis := NewGenericVisitor(func(x any) bool { - switch x := x.(type) { + switch y := x.(type) { case Call: return f(x) case *Expr: - if x.IsCall() { + if y.IsCall() { return f(x) } case *Head: // GenericVisitor doesn't walk the rule head ref - walkCalls(x.Reference, f) + walkCalls(y.Reference, f) } return false }) diff --git a/v1/ast/strings.go b/v1/ast/strings.go index 8447522412..72ec03f8cc 100644 --- a/v1/ast/strings.go +++ b/v1/ast/strings.go @@ -48,6 +48,8 @@ func ValueName(x Value) string { return "objectcomprehension" case *SetComprehension: return "setcomprehension" + case *TemplateString: + return "templatestring" } return TypeName(x) diff --git a/v1/ast/term_bench_test.go b/v1/ast/term_bench_test.go index 6d63213437..acb5b14264 100644 --- a/v1/ast/term_bench_test.go +++ b/v1/ast/term_bench_test.go @@ -274,6 +274,35 @@ func BenchmarkObjectCreationAndLookup(b *testing.B) { } } +// insert 38148 30049 ns/op 58912 B/op 528 allocs/op +// terms_array 65698 17079 ns/op 34680 B/op 506 allocs/op +func BenchmarkObjectCreateWithInsertVsTermsArray(b *testing.B) { + n := 500 + interned := make([]*Term, n) + for i := range n { + interned[i] = InternedTerm(i) + } + + b.Run("insert", func(b *testing.B) { + for b.Loop() { + obj := NewObject() + for i := range n { + obj.Insert(interned[i], interned[i]) + } + } + }) + b.Run("terms array", func(b *testing.B) { + for b.Loop() { + terms := make([][2]*Term, n) + for i := range n { + terms[i][0] = interned[i] + terms[i][1] = interned[i] + } + _ = NewObject(terms...) + } + }) +} + func BenchmarkLazyObjectLookup(b *testing.B) { sizes := []int{5, 50, 500, 5000} for _, n := range sizes { @@ -783,6 +812,24 @@ func BenchmarkInterfaceToValueInt(b *testing.B) { }) } +func BenchmarkValueToInterfaceInt(b *testing.B) { + term := MustParseTerm(`{ + "foo": [1, "two", true, false, null, {3: 4}], + "bar": 5.67, + "baz": {"a": "b", "c": "d"}, + "set": {10, 20, 30} + }`) + + opt := JSONOpt{SortSets: true, CopyMaps: true} + + for b.Loop() { + _, err := valueToInterface(term.Value, nil, opt) + if err != nil { + b.Fatal(err) + } + } +} + // without_conflict 258.8 ns/op 440 B/op 11 allocs/op // use NewObject // with_conflict 290.6 ns/op 440 B/op 11 allocs/op // // without_conflict 220.5 ns/op 408 B/op 9 allocs/op // use newobject with size diff --git a/v1/ast/transform.go b/v1/ast/transform.go index 7d0a75bbd5..a71bc0a77c 100644 --- a/v1/ast/transform.go +++ b/v1/ast/transform.go @@ -303,29 +303,29 @@ func Transform(t Transformer, x any) (any, error) { // TransformRefs calls the function f on all references under x. func TransformRefs(x any, f func(Ref) (Value, error)) (any, error) { - t := &GenericTransformer{func(x any) (any, error) { + t := NewGenericTransformer(func(x any) (any, error) { if r, ok := x.(Ref); ok { return f(r) } return x, nil - }} + }) return Transform(t, x) } // TransformVars calls the function f on all vars under x. func TransformVars(x any, f func(Var) (Value, error)) (any, error) { - t := &GenericTransformer{func(x any) (any, error) { + t := NewGenericTransformer(func(x any) (any, error) { if v, ok := x.(Var); ok { return f(v) } return x, nil - }} + }) return Transform(t, x) } -// TransformComprehensions calls the functio nf on all comprehensions under x. +// TransformComprehensions calls the function f on all comprehensions under x. func TransformComprehensions(x any, f func(any) (Value, error)) (any, error) { - t := &GenericTransformer{func(x any) (any, error) { + t := NewGenericTransformer(func(x any) (any, error) { switch x := x.(type) { case *ArrayComprehension: return f(x) @@ -335,7 +335,7 @@ func TransformComprehensions(x any, f func(any) (Value, error)) (any, error) { return f(x) } return x, nil - }} + }) return Transform(t, x) } diff --git a/v1/bundle/store.go b/v1/bundle/store.go index e79aea5200..f6c13c75fc 100644 --- a/v1/bundle/store.go +++ b/v1/bundle/store.go @@ -970,7 +970,7 @@ func compileModules(compiler *ast.Compiler, m metrics.Metrics, bundles map[strin m.Timer(metrics.RegoModuleCompile).Start() defer m.Timer(metrics.RegoModuleCompile).Stop() - modules := map[string]*ast.Module{} + modules := make(map[string]*ast.Module, len(compiler.Modules)+len(extraModules)+len(bundles)) // preserve any modules already on the compiler maps.Copy(modules, compiler.Modules) diff --git a/v1/format/format.go b/v1/format/format.go index 69b6ec51ee..251f45c3a6 100644 --- a/v1/format/format.go +++ b/v1/format/format.go @@ -27,8 +27,6 @@ import ( const defaultLocationFile = "__format_default__" var ( - elseVar ast.Value = ast.Var("else") - expandedConst = ast.NewBody(ast.NewExpr(ast.InternedTerm(true))) commentsSlicePool = util.NewSlicePool[*ast.Comment](50) varRegexp = regexp.MustCompile("^[[:alpha:]_][[:alpha:][:digit:]_]*$") @@ -732,7 +730,7 @@ func (w *writer) writeElse(rule *ast.Rule, comments []*ast.Comment) ([]*ast.Comm rule.Else.Head.Name = "else" // NOTE(sr): whaaat - elseHeadReference := ast.NewTerm(elseVar) // construct a reference for the term + elseHeadReference := ast.VarTerm("else") // construct a reference for the term elseHeadReference.Location = rule.Else.Head.Location // and set the location to match the rule location rule.Else.Head.Reference = ast.Ref{elseHeadReference} diff --git a/v1/repl/repl.go b/v1/repl/repl.go index 3f65a593b7..f84a6f7475 100644 --- a/v1/repl/repl.go +++ b/v1/repl/repl.go @@ -1293,9 +1293,8 @@ func (r *REPL) loadModules(ctx context.Context, txn storage.Transaction) (map[st } func (r *REPL) printTypes(_ context.Context, typeEnv *ast.TypeEnv, body ast.Body) { - ast.WalkRefs(body, func(ref ast.Ref) bool { - fmt.Fprintf(r.output, "# %v: %v\n", ref, typeEnv.Get(ref)) + fmt.Fprintf(r.output, "# %v: %v\n", ref, typeEnv.GetByRef(ref)) return false }) diff --git a/v1/topdown/print.go b/v1/topdown/print.go index 800ca795ad..acf55f0f3f 100644 --- a/v1/topdown/print.go +++ b/v1/topdown/print.go @@ -28,7 +28,6 @@ func (h printHook) Print(_ print.Context, msg string) error { } func builtinPrint(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error { - if bctx.PrintHook == nil { return iter(nil) } @@ -40,7 +39,7 @@ func builtinPrint(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term buf := make([]string, arr.Len()) - err = builtinPrintCrossProductOperands(bctx, buf, arr, 0, func(buf []string) error { + err = builtinPrintCrossProductOperands(bctx.Location, buf, arr, 0, func(buf []string) error { pctx := print.Context{ Context: bctx.Context, Location: bctx.Location, @@ -54,31 +53,32 @@ func builtinPrint(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term return iter(nil) } -func builtinPrintCrossProductOperands(bctx BuiltinContext, buf []string, operands *ast.Array, i int, f func([]string) error) error { - +func builtinPrintCrossProductOperands(loc *ast.Location, buf []string, operands *ast.Array, i int, f func([]string) error) error { if i >= operands.Len() { return f(buf) } + operand := operands.Elem(i) + // We allow primitives ... - switch x := operands.Elem(i).Value.(type) { + switch x := operand.Value.(type) { case ast.String: buf[i] = string(x) - return builtinPrintCrossProductOperands(bctx, buf, operands, i+1, f) + return builtinPrintCrossProductOperands(loc, buf, operands, i+1, f) case ast.Number, ast.Boolean, ast.Null: buf[i] = x.String() - return builtinPrintCrossProductOperands(bctx, buf, operands, i+1, f) + return builtinPrintCrossProductOperands(loc, buf, operands, i+1, f) } // ... but all other operand types must be sets. - xs, ok := operands.Elem(i).Value.(ast.Set) + xs, ok := operand.Value.(ast.Set) if !ok { - return Halt{Err: internalErr(bctx.Location, fmt.Sprintf("illegal argument type: %v", ast.ValueName(operands.Elem(i).Value)))} + return Halt{Err: internalErr(loc, "illegal argument type: "+ast.ValueName(operand.Value))} } if xs.Len() == 0 { buf[i] = "" - return builtinPrintCrossProductOperands(bctx, buf, operands, i+1, f) + return builtinPrintCrossProductOperands(loc, buf, operands, i+1, f) } return xs.Iter(func(x *ast.Term) error { @@ -88,7 +88,7 @@ func builtinPrintCrossProductOperands(bctx BuiltinContext, buf []string, operand default: buf[i] = v.String() } - return builtinPrintCrossProductOperands(bctx, buf, operands, i+1, f) + return builtinPrintCrossProductOperands(loc, buf, operands, i+1, f) }) } diff --git a/v1/topdown/query.go b/v1/topdown/query.go index f81402eb32..a65f8a312c 100644 --- a/v1/topdown/query.go +++ b/v1/topdown/query.go @@ -134,7 +134,7 @@ func (q *Query) WithTracer(tracer Tracer) *Query { // WithQueryTracer adds a query tracer to use during evaluation. This is optional. // Disabled QueryTracers will be ignored. func (q *Query) WithQueryTracer(tracer QueryTracer) *Query { - if !tracer.Enabled() { + if tracer == nil || !tracer.Enabled() { return q } diff --git a/v1/topdown/template_string.go b/v1/topdown/template_string.go index 4df0d12acf..0e705b81bf 100644 --- a/v1/topdown/template_string.go +++ b/v1/topdown/template_string.go @@ -20,7 +20,7 @@ func builtinTemplateString(bctx BuiltinContext, operands []*ast.Term, iter func( buf := make([]string, arr.Len()) var count int - err = builtinPrintCrossProductOperands(bctx, buf, arr, 0, func(buf []string) error { + err = builtinPrintCrossProductOperands(bctx.Location, buf, arr, 0, func(buf []string) error { count += 1 // Precautionary run-time assertion that template-strings can't produce multiple outputs; e.g. for custom relation type built-ins not known at compile-time. if count > 1 { @@ -37,8 +37,7 @@ func builtinTemplateString(bctx BuiltinContext, operands []*ast.Term, iter func( return err } - str := ast.StringTerm(strings.Join(buf, "")) - return iter(str) + return iter(ast.StringTerm(strings.Join(buf, ""))) } func init() { diff --git a/v1/topdown/template_string_test.go b/v1/topdown/template_string_test.go index 751bf1a688..578600e437 100644 --- a/v1/topdown/template_string_test.go +++ b/v1/topdown/template_string_test.go @@ -7,52 +7,52 @@ import ( "github.com/open-policy-agent/opa/v1/ast" ) -func TestBuiltinTemplateString(t *testing.T) { - tests := []struct { - note string - parts *ast.Array - expRes *ast.Term - expErr string - }{ - { - note: "no parts", - parts: ast.NewArray(), - expRes: ast.StringTerm(""), - }, - { - note: "single string part", - parts: ast.NewArray(ast.StringTerm("foo")), - expRes: ast.StringTerm("foo"), - }, - { - note: "single undefined part", - parts: ast.NewArray(ast.SetTerm()), - expRes: ast.StringTerm(""), - }, - { - note: "primitives", - parts: ast.NewArray(ast.StringTerm("foo"), ast.NumberTerm("42"), ast.BooleanTerm(false), ast.NullTerm()), - expRes: ast.StringTerm("foo42falsenull"), - }, - { - note: "collections", - parts: ast.NewArray( - ast.SetTerm(ast.ArrayTerm()), ast.StringTerm(" "), - ast.SetTerm(ast.ArrayTerm(ast.StringTerm("a"), ast.StringTerm("b"))), ast.StringTerm(" "), - ast.SetTerm(ast.SetTerm()), ast.StringTerm(" "), - ast.SetTerm(ast.SetTerm(ast.StringTerm("c"))), ast.StringTerm(" "), - ast.SetTerm(ast.ObjectTerm()), ast.StringTerm(" "), - ast.SetTerm(ast.ObjectTerm(ast.Item(ast.StringTerm("d"), ast.StringTerm("e")))), - ), - expRes: ast.StringTerm(`[] ["a", "b"] set() {"c"} {} {"d": "e"}`), - }, - { - note: "multiple outputs", - parts: ast.NewArray(ast.SetTerm(ast.BooleanTerm(true), ast.BooleanTerm(false))), - expErr: "eval_conflict_error: template-strings must not produce multiple outputs", - }, - } +var tests = []struct { + note string + parts *ast.Array + expRes *ast.Term + expErr string +}{ + { + note: "no parts", + parts: ast.NewArray(), + expRes: ast.StringTerm(""), + }, + { + note: "single string part", + parts: ast.NewArray(ast.StringTerm("foo")), + expRes: ast.StringTerm("foo"), + }, + { + note: "single undefined part", + parts: ast.NewArray(ast.SetTerm()), + expRes: ast.StringTerm(""), + }, + { + note: "primitives", + parts: ast.NewArray(ast.StringTerm("foo"), ast.NumberTerm("42"), ast.BooleanTerm(false), ast.NullTerm()), + expRes: ast.StringTerm("foo42falsenull"), + }, + { + note: "collections", + parts: ast.NewArray( + ast.SetTerm(ast.ArrayTerm()), ast.StringTerm(" "), + ast.SetTerm(ast.ArrayTerm(ast.StringTerm("a"), ast.StringTerm("b"))), ast.StringTerm(" "), + ast.SetTerm(ast.SetTerm()), ast.StringTerm(" "), + ast.SetTerm(ast.SetTerm(ast.StringTerm("c"))), ast.StringTerm(" "), + ast.SetTerm(ast.ObjectTerm()), ast.StringTerm(" "), + ast.SetTerm(ast.ObjectTerm(ast.Item(ast.StringTerm("d"), ast.StringTerm("e")))), + ), + expRes: ast.StringTerm(`[] ["a", "b"] set() {"c"} {} {"d": "e"}`), + }, + { + note: "multiple outputs", + parts: ast.NewArray(ast.SetTerm(ast.BooleanTerm(true), ast.BooleanTerm(false))), + expErr: "eval_conflict_error: template-strings must not produce multiple outputs", + }, +} +func TestBuiltinTemplateString(t *testing.T) { for _, tc := range tests { t.Run(tc.note, func(t *testing.T) { var result *ast.Term @@ -82,3 +82,23 @@ func TestBuiltinTemplateString(t *testing.T) { }) } } + +// BenchmarkBuiltinTemplateString/no_parts-16 13434396 82.17 ns/op 344 B/op 4 allocs/op +// BenchmarkBuiltinTemplateString/single_string_part-16 11506334 106.0 ns/op 376 B/op 6 allocs/op +// BenchmarkBuiltinTemplateString/single_undefined_part-16 11367075 106.0 ns/op 376 B/op 6 allocs/op +// BenchmarkBuiltinTemplateString/primitives-16 8217890 144.9 ns/op 440 B/op 7 allocs/op +// BenchmarkBuiltinTemplateString/collections-16 2056494 583.7 ns/op 1144 B/op 28 allocs/op +// BenchmarkBuiltinTemplateString/multiple_outputs-16 9424003 128.8 ns/op 480 B/op 7 allocs/op +func BenchmarkBuiltinTemplateString(b *testing.B) { + for _, tc := range tests { + b.Run(tc.note, func(b *testing.B) { + bctx := BuiltinContext{} + oper := []*ast.Term{ast.NewTerm(tc.parts)} + iter := eqIter(tc.expRes) + + for b.Loop() { + _ = builtinTemplateString(bctx, oper, iter) + } + }) + } +} diff --git a/v1/topdown/topdown_bench_test.go b/v1/topdown/topdown_bench_test.go index 0934c6f7e1..0818d3d6be 100644 --- a/v1/topdown/topdown_bench_test.go +++ b/v1/topdown/topdown_bench_test.go @@ -969,3 +969,66 @@ func BenchmarkObjectGetFromBaseDoc(b *testing.B) { } } } + +// templatestring 216386 5524 ns/op 10074 B/op 158 allocs/op +// templatestring 5219 ns/op 8337 B/op 152 allocs/op don't pass bctx +// -- +// concat 272054 4500 ns/op 6728 B/op 123 allocs/op +// sprintf 280381 4421 ns/op 6365 B/op 121 allocs/op +func BenchmarkTemplateStringVsConcatVsSprintf(b *testing.B) { + ctx := b.Context() + store := inmem.NewFromObject(map[string]any{}) + + modBase := `package test + + foo := "foo" + bar := "bar" + baz := "baz" + + ` + tests := []struct { + name string + snippet string + }{ + { + name: "templatestring", + snippet: `s := $"{foo}-{bar}-{baz}"`, + }, + { + name: "concat", + snippet: `s := concat("-", [foo, bar, baz])`, + }, + { + name: "sprintf", + snippet: `s := sprintf("%s-%s-%s", [foo, bar, baz])`, + }, + } + + for _, tc := range tests { + policy := modBase + tc.snippet + + query := ast.MustParseBody("data.test.s = s") + compiler := ast.MustCompileModules(map[string]string{"test.rego": policy}) + + b.Run(tc.name, func(b *testing.B) { + err := storage.Txn(ctx, store, storage.TransactionParams{}, func(txn storage.Transaction) error { + q := NewQuery(query).WithCompiler(compiler).WithStore(store).WithTransaction(txn) + for b.Loop() { + rs, err := q.Run(ctx) + if err != nil { + return err + } + result := rs[0][ast.Var("s")] + expected := ast.String("foo-bar-baz") + if !result.Value.(ast.String).Equal(expected) { + b.Fatalf("unexpected result: %v (expected: %v)", result, expected) + } + } + return nil + }) + if err != nil { + b.Fatal(err) + } + }) + } +}