ewilliams28/releases
1
0
Fork 0
mirror of https://github.com/open-policy-agent/opa.git synced 2026-08-12 19:32:48 -06:00
Code Issues Packages Projects Releases Wiki Activity

Template string performance improvements and more (#8143)

Browse Source 
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 <anders.eknert@apple.com>
This commit is contained in:
Anders Eknert
2025-12-18 12:04:18 +01:00
committed by GitHub
parent f5c3743eed
commit 5a0dc476be
17 changed files with 297 additions and 211 deletions
Download Patch File Download Diff File Expand all files Collapse all files
v1/ast/annotations.go
+1 -12
View File
@@ -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 {
v1/ast/check.go
+50 -70
View File
@@ -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
}
v1/ast/env.go
+24 -41
View File
@@ -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
}
v1/ast/interning.go
+11 -4
View File
@@ -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"),
}
v1/ast/policy.go
+1 -1
View File
@@ -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) {
v1/ast/rego_v1.go
+9 -10
View File
@@ -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
})
v1/ast/strings.go
+2
View File
@@ -48,6 +48,8 @@ func ValueName(x Value) string {
return "objectcomprehension"
case *SetComprehension:
return "setcomprehension"
case *TemplateString:
return "templatestring"
}
return TypeName(x)
v1/ast/term_bench_test.go
+47
View File
@@ -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
v1/ast/transform.go
+7 -7
View File
@@ -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)
}
v1/bundle/store.go
+1 -1
View File
@@ -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)
v1/format/format.go
+1 -3
View File
@@ -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}
v1/repl/repl.go
+1 -2
View File
@@ -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
})
v1/topdown/print.go
+11 -11
View File
@@ -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] = "<undefined>"
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)
})
}
v1/topdown/query.go
+1 -1
View File
@@ -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
}
v1/topdown/template_string.go
+2 -3
View File
@@ -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() {
v1/topdown/template_string_test.go
+65 -45
View File
@@ -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("<undefined>"),
},
{
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("<undefined>"),
},
{
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)
}
})
}
}
v1/topdown/topdown_bench_test.go
+63
View File
@@ -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)
}
})
}
}