compile+types: Speed up typechecker when working with Refs (#5307)

* compile/compile_bench_test: Add basic compilation benchmark.
* types/types: Replace linear copy loops with copy() Go builtin.
* types/types: Switch linear scan -> binary search.

This commit switches out the linear scan in the `Any` type's `Contains`
method for a more efficient binary search. This results in around a 3-5x
speedup for policy compilation as a whole, according to the benchmark.

* ast/visit: Remove for loop item copies.

This commit refactors the visitor functions in the `ast/visit` package
to avoid extra copying of items on each loop iteration. This shaves off
around 10-15% memory usage during type-checking during compilation, and
provides speed benefits due to reduced GC as a result.

Signed-off-by: Philip Conrad <philipaconrad@gmail.com>
This commit is contained in:
Philip Conrad
2022-10-27 03:48:39 -04:00
committed by GitHub
parent da8afed08d
commit 3e640489bd
3 changed files with 215 additions and 98 deletions
+93 -86
View File
@@ -46,17 +46,17 @@ func walk(v Visitor, x interface{}) {
switch x := x.(type) {
case *Module:
Walk(w, x.Package)
for _, i := range x.Imports {
Walk(w, i)
for i := range x.Imports {
Walk(w, x.Imports[i])
}
for _, r := range x.Rules {
Walk(w, r)
for i := range x.Rules {
Walk(w, x.Rules[i])
}
for _, a := range x.Annotations {
Walk(w, a)
for i := range x.Annotations {
Walk(w, x.Annotations[i])
}
for _, c := range x.Comments {
Walk(w, c)
for i := range x.Comments {
Walk(w, x.Comments[i])
}
case *Package:
Walk(w, x.Path)
@@ -79,20 +79,20 @@ func walk(v Visitor, x interface{}) {
Walk(w, x.Value)
}
case Body:
for _, e := range x {
Walk(w, e)
for i := range x {
Walk(w, x[i])
}
case Args:
for _, t := range x {
Walk(w, t)
for i := range x {
Walk(w, x[i])
}
case *Expr:
switch ts := x.Terms.(type) {
case *Term, *SomeDecl, *Every:
Walk(w, ts)
case []*Term:
for _, t := range ts {
Walk(w, t)
for i := range ts {
Walk(w, ts[i])
}
}
for i := range x.With {
@@ -104,8 +104,8 @@ func walk(v Visitor, x interface{}) {
case *Term:
Walk(w, x.Value)
case Ref:
for _, t := range x {
Walk(w, t)
for i := range x {
Walk(w, x[i])
}
case *object:
x.Foreach(func(k, vv *Term) {
@@ -131,8 +131,8 @@ func walk(v Visitor, x interface{}) {
Walk(w, x.Term)
Walk(w, x.Body)
case Call:
for _, t := range x {
Walk(w, t)
for i := range x {
Walk(w, x[i])
}
case *Every:
if x.Key != nil {
@@ -282,17 +282,17 @@ func (vis *GenericVisitor) Walk(x interface{}) {
switch x := x.(type) {
case *Module:
vis.Walk(x.Package)
for _, i := range x.Imports {
vis.Walk(i)
for i := range x.Imports {
vis.Walk(x.Imports[i])
}
for _, r := range x.Rules {
vis.Walk(r)
for i := range x.Rules {
vis.Walk(x.Rules[i])
}
for _, a := range x.Annotations {
vis.Walk(a)
for i := range x.Annotations {
vis.Walk(x.Annotations[i])
}
for _, c := range x.Comments {
vis.Walk(c)
for i := range x.Comments {
vis.Walk(x.Comments[i])
}
case *Package:
vis.Walk(x.Path)
@@ -315,20 +315,20 @@ func (vis *GenericVisitor) Walk(x interface{}) {
vis.Walk(x.Value)
}
case Body:
for _, e := range x {
vis.Walk(e)
for i := range x {
vis.Walk(x[i])
}
case Args:
for _, t := range x {
vis.Walk(t)
for i := range x {
vis.Walk(x[i])
}
case *Expr:
switch ts := x.Terms.(type) {
case *Term, *SomeDecl, *Every:
vis.Walk(ts)
case []*Term:
for _, t := range ts {
vis.Walk(t)
for i := range ts {
vis.Walk(ts[i])
}
}
for i := range x.With {
@@ -340,8 +340,8 @@ func (vis *GenericVisitor) Walk(x interface{}) {
case *Term:
vis.Walk(x.Value)
case Ref:
for _, t := range x {
vis.Walk(t)
for i := range x {
vis.Walk(x[i])
}
case *object:
x.Foreach(func(k, v *Term) {
@@ -353,8 +353,9 @@ func (vis *GenericVisitor) Walk(x interface{}) {
vis.Walk(t)
})
case Set:
for _, t := range x.Slice() {
vis.Walk(t)
xSlice := x.Slice()
for i := range xSlice {
vis.Walk(xSlice[i])
}
case *ArrayComprehension:
vis.Walk(x.Term)
@@ -367,8 +368,8 @@ func (vis *GenericVisitor) Walk(x interface{}) {
vis.Walk(x.Term)
vis.Walk(x.Body)
case Call:
for _, t := range x {
vis.Walk(t)
for i := range x {
vis.Walk(x[i])
}
case *Every:
if x.Key != nil {
@@ -408,17 +409,17 @@ func (vis *BeforeAfterVisitor) Walk(x interface{}) {
switch x := x.(type) {
case *Module:
vis.Walk(x.Package)
for _, i := range x.Imports {
vis.Walk(i)
for i := range x.Imports {
vis.Walk(x.Imports[i])
}
for _, r := range x.Rules {
vis.Walk(r)
for i := range x.Rules {
vis.Walk(x.Rules[i])
}
for _, a := range x.Annotations {
vis.Walk(a)
for i := range x.Annotations {
vis.Walk(x.Annotations[i])
}
for _, c := range x.Comments {
vis.Walk(c)
for i := range x.Comments {
vis.Walk(x.Comments[i])
}
case *Package:
vis.Walk(x.Path)
@@ -445,20 +446,20 @@ func (vis *BeforeAfterVisitor) Walk(x interface{}) {
vis.Walk(x.Value)
}
case Body:
for _, e := range x {
vis.Walk(e)
for i := range x {
vis.Walk(x[i])
}
case Args:
for _, t := range x {
vis.Walk(t)
for i := range x {
vis.Walk(x[i])
}
case *Expr:
switch ts := x.Terms.(type) {
case *Term, *SomeDecl, *Every:
vis.Walk(ts)
case []*Term:
for _, t := range ts {
vis.Walk(t)
for i := range ts {
vis.Walk(ts[i])
}
}
for i := range x.With {
@@ -470,8 +471,8 @@ func (vis *BeforeAfterVisitor) Walk(x interface{}) {
case *Term:
vis.Walk(x.Value)
case Ref:
for _, t := range x {
vis.Walk(t)
for i := range x {
vis.Walk(x[i])
}
case *object:
x.Foreach(func(k, v *Term) {
@@ -483,8 +484,9 @@ func (vis *BeforeAfterVisitor) Walk(x interface{}) {
vis.Walk(t)
})
case Set:
for _, t := range x.Slice() {
vis.Walk(t)
xSlice := x.Slice()
for i := range xSlice {
vis.Walk(xSlice[i])
}
case *ArrayComprehension:
vis.Walk(x.Term)
@@ -497,8 +499,8 @@ func (vis *BeforeAfterVisitor) Walk(x interface{}) {
vis.Walk(x.Term)
vis.Walk(x.Body)
case Call:
for _, t := range x {
vis.Walk(t)
for i := range x {
vis.Walk(x[i])
}
case *Every:
if x.Key != nil {
@@ -559,8 +561,9 @@ func (vis *VarVisitor) visit(v interface{}) bool {
}
if vis.params.SkipRefHead {
if r, ok := v.(Ref); ok {
for _, t := range r[1:] {
vis.Walk(t)
rSlice := r[1:]
for i := range rSlice {
vis.Walk(rSlice[i])
}
return true
}
@@ -592,14 +595,15 @@ func (vis *VarVisitor) visit(v interface{}) bool {
switch v := v.(type) {
case *Expr:
if terms, ok := v.Terms.([]*Term); ok {
for _, t := range terms[0].Value.(Ref)[1:] {
vis.Walk(t)
termSlice := terms[0].Value.(Ref)[1:]
for i := range termSlice {
vis.Walk(termSlice[i])
}
for i := 1; i < len(terms); i++ {
vis.Walk(terms[i])
}
for _, w := range v.With {
vis.Walk(w)
for i := range v.With {
vis.Walk(v.With[i])
}
return true
}
@@ -614,13 +618,15 @@ func (vis *VarVisitor) visit(v interface{}) bool {
return true
case *With:
if ref, ok := v.Target.Value.(Ref); ok {
for _, t := range ref[1:] {
vis.Walk(t)
refSlice := ref[1:]
for i := range refSlice {
vis.Walk(refSlice[i])
}
}
if ref, ok := v.Value.Value.(Ref); ok {
for _, t := range ref[1:] {
vis.Walk(t)
refSlice := ref[1:]
for i := range refSlice {
vis.Walk(refSlice[i])
}
} else {
vis.Walk(v.Value)
@@ -645,14 +651,14 @@ func (vis *VarVisitor) Walk(x interface{}) {
switch x := x.(type) {
case *Module:
vis.Walk(x.Package)
for _, i := range x.Imports {
vis.Walk(i)
for i := range x.Imports {
vis.Walk(x.Imports[i])
}
for _, r := range x.Rules {
vis.Walk(r)
for i := range x.Rules {
vis.Walk(x.Rules[i])
}
for _, c := range x.Comments {
vis.Walk(c)
for i := range x.Comments {
vis.Walk(x.Comments[i])
}
case *Package:
vis.Walk(x.Path)
@@ -680,20 +686,20 @@ func (vis *VarVisitor) Walk(x interface{}) {
vis.Walk(x.Value)
}
case Body:
for _, e := range x {
vis.Walk(e)
for i := range x {
vis.Walk(x[i])
}
case Args:
for _, t := range x {
vis.Walk(t)
for i := range x {
vis.Walk(x[i])
}
case *Expr:
switch ts := x.Terms.(type) {
case *Term, *SomeDecl, *Every:
vis.Walk(ts)
case []*Term:
for _, t := range ts {
vis.Walk(t)
for i := range ts {
vis.Walk(ts[i])
}
}
for i := range x.With {
@@ -705,8 +711,8 @@ func (vis *VarVisitor) Walk(x interface{}) {
case *Term:
vis.Walk(x.Value)
case Ref:
for _, t := range x {
vis.Walk(t)
for i := range x {
vis.Walk(x[i])
}
case *object:
x.Foreach(func(k, v *Term) {
@@ -718,8 +724,9 @@ func (vis *VarVisitor) Walk(x interface{}) {
vis.Walk(t)
})
case Set:
for _, t := range x.Slice() {
vis.Walk(t)
xSlice := x.Slice()
for i := range xSlice {
vis.Walk(xSlice[i])
}
case *ArrayComprehension:
vis.Walk(x.Term)
@@ -732,8 +739,8 @@ func (vis *VarVisitor) Walk(x interface{}) {
vis.Walk(x.Term)
vis.Walk(x.Body)
case Call:
for _, t := range x {
vis.Walk(t)
for i := range x {
vis.Walk(x[i])
}
case *Every:
if x.Key != nil {
+110
View File
@@ -0,0 +1,110 @@
package compile
import (
"context"
"fmt"
"testing"
"github.com/open-policy-agent/opa/util/test"
)
// type compileBenchTestData struct {
// filename string
// module string
// }
func BenchmarkCompileDynamicPolicy(b *testing.B) {
// This benchmarks the compiler against increasingly large numbers of dynamically-selected policies.
// See: https://github.com/open-policy-agent/opa/issues/5216
//ctx := context.Background()
numPolicies := []int{1000, 2500, 5000, 7500, 10000}
testcases := map[int]map[string]string{}
for _, n := range numPolicies {
testcases[n] = generateDynamicPolicyBenchmarkData(n)
}
b.ResetTimer()
for _, n := range numPolicies {
b.Run(fmt.Sprintf("%d", n), func(b *testing.B) {
test.WithTempFS(testcases[n], func(root string) {
b.ResetTimer()
compiler := New().
WithPaths(root)
err := compiler.Build(context.Background())
if err != nil {
b.Fatal("unexpected error", err)
}
})
// store := inmem.NewFromObject(map[string]interface{}{"objs": generateMockPolicy(n)})
// module := `package test
// combined := {k: true | s := data.objs[_]; s[k]}`
// query := ast.MustParseBody("data.test.combined")
// compiler := ast.MustCompileModules(map[string]string{
// "test.rego": module,
// })
// b.ResetTimer()
// for i := 0; i < b.N; i++ {
// err := storage.Txn(ctx, store, storage.TransactionParams{}, func(txn storage.Transaction) error {
// q := NewQuery(query).
// WithCompiler(compiler).
// WithStore(store).
// WithTransaction(txn)
// _, err := q.Run(ctx)
// if err != nil {
// return err
// }
// return nil
// })
// if err != nil {
// b.Fatal(err)
// }
// }
})
}
}
func generateDynamicPolicyBenchmarkData(N int) map[string]string {
files := map[string]string{
"main.rego": `
package main
denies[x] {
x := data.policies[input.type][input.subtype][_].denies[_]
}
any_denies {
denies[_]
}
allow {
not any_denies
}`,
}
for i := 0; i < N; i++ {
files[fmt.Sprintf("policy%d.rego", i)] = generateMockPolicy(i)
}
return files
}
func generateMockPolicy(N int) string {
return fmt.Sprintf(`package policies["%d"]["%d"].policy%d
denies[x] {
input.attribute == "%d"
x := "policy%d"
}`, N, N, N, N, N)
}
+12 -12
View File
@@ -430,9 +430,7 @@ var A = NewAny()
// NewAny returns a new Any type.
func NewAny(of ...Type) Any {
sl := make(Any, len(of))
for i := range sl {
sl[i] = of[i]
}
copy(sl, of)
sort.Sort(typeSlice(sl))
return sl
}
@@ -442,10 +440,14 @@ func (t Any) Contains(other Type) bool {
if _, ok := other.(*Function); ok {
return false
}
for i := range t {
if Compare(t[i], other) == 0 {
return true
}
// Note(philipc): We used to do this as a linear search.
// Since this is always sorted, we can use a binary search instead.
i := sort.Search(len(t), func(i int) bool {
return Compare(t[i], other) >= 0
})
if i < len(t) && Compare(t[i], other) == 0 {
// x is present at t[i]
return true
}
return len(t) == 0
}
@@ -492,9 +494,7 @@ func (t Any) Union(other Any) Any {
return other
}
cpy := make(Any, len(t))
for i := range cpy {
cpy[i] = t[i]
}
copy(cpy, t)
for i := range other {
if !cpy.Contains(other[i]) {
cpy = append(cpy, other[i])
@@ -929,8 +929,8 @@ func Values(a Type) Type {
return Or(tpe, a.dynamic)
case *Object:
var tpe Type
for _, v := range a.static {
tpe = Or(tpe, v.Value)
for i := range a.static {
tpe = Or(tpe, a.static[i].Value)
}
if a.dynamic != nil {
tpe = Or(tpe, a.dynamic.Value)