Skip to content

Commit

Permalink
wasm: have planner optimize case without rules in leaves (do nothing) (…
Browse files Browse the repository at this point in the history 
…#3089)

In our previous optimization change, #3058, we've been falling back to the
old behaviour (with introduces cross-products) whenever our optimization
didn't seem to apply.

This PR extends that a little: when there are only leaves (matching our
optimization lookup) that do not have rules (i.e. empty packages), we decide
that nothing package-related is to be done at all. So, we'll go on looking at
base data only.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>
  • Loading branch information
@srenatus
srenatus authored Jan 25, 2021
1 parent 1d96675 commit dcdbf36
Show file tree
Hide file tree
Showing 4 changed files with 60 additions and 31 deletions.
25 changes: 18 additions & 7 deletions internal/planner/planner.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1526,9 +1526,15 @@ func (p *Planner) planRefData(virtual *ruletrie, base *baseptr, ref ast.Ref, ind
// CallDynamicStatement
// NOTE(sr): we do it on the first index because later on, the recursion
// on subtrees of virtual already lost parts of the path we've taken.
if index == 1 {
if index == 1 && virtual != nil {
rulesets, stmts, locals, index, optimize := p.optimizeLookup(virtual, ref)
if optimize {
// If there are no rulesets in a situation that otherwise would
// allow for a call_indirect optimization, then there's nothing
// to do for this ref, except scanning the base document.
if len(rulesets) == 0 {
return p.planRefData(nil, base, ref, 1, iter) // ignore index returned by optimizeLookup
}
// plan rules
for _, rules := range rulesets {
if _, err := p.planRules(rules); err != nil {
Expand Down Expand Up @@ -1650,11 +1656,14 @@ func (p *Planner) planRefData(virtual *ruletrie, base *baseptr, ref ast.Ref, ind
}

p.ltarget = base.local
return p.planRefDataBaseScan(base.path, ref, index, lexclude, iter)
}

// Perform a scan of the base documents starting from the location referred
// to by the data pointer. Use the set we built above to avoid revisiting
// sub trees.
return p.planRefRec(base.path, 0, func() error {
// Perform a scan of the base documents starting from the location referred
// to by the 'path' data pointer. Use the set (planned into 'lexclude') to
// avoid revisiting sub trees.
func (p *Planner) planRefDataBaseScan(path ast.Ref, ref ast.Ref, index int, lexclude *ir.Local, iter planiter) error {
return p.planRefRec(path, 0, func() error {
return p.planScan(ref[index], func(lkey ir.Local) error {
if lexclude != nil {
lignore := p.newLocal()
Expand Down Expand Up @@ -2072,6 +2081,7 @@ func (p *Planner) optimizeLookup(t *ruletrie, ref ast.Ref) ([][]*ast.Rule, []ir.
all = all && len(node.Children()) == 0
}
if all {
p.debug(fmt.Sprintf("ref %s: all nodes have 0 children, break", ref[0:index+1]))
break
}
}
Expand All @@ -2081,7 +2091,7 @@ func (p *Planner) optimizeLookup(t *ruletrie, ref ast.Ref) ([][]*ast.Rule, []ir.
// if there hasn't been any var, we're not making things better by
// introducing CallDynamicStmt
if !opt {
return dont("no vars at all")
return dont("no vars seen before trie descend encountered no children")
}

for _, node := range nodes {
Expand All @@ -2096,7 +2106,8 @@ func (p *Planner) optimizeLookup(t *ruletrie, ref ast.Ref) ([][]*ast.Rule, []ir.
}
}
if len(res) == 0 {
return dont("no rule leaves")
p.debug(fmt.Sprintf("ref %s: nothing to plan, no rule leaves", ref[0:index+1]))
return nil, nil, nil, index, true
}

// If we've made it here, optimization is possible -- let's plan strings!
Expand Down
15 changes: 13 additions & 2 deletions internal/planner/planner_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -317,6 +317,11 @@ p = 1
q = 2`,
},
},
{
note: "cross product with non-ground refs to packages, 'no rules leaves' case",
queries: []string{`x := "aaa"; y := data.foo[x].bar.baz`},
modules: []string{`package foo.aaa.bar`, `package foo.bbb.bar`},
},
}

for _, tc := range tests {
Expand Down Expand Up @@ -346,6 +351,9 @@ q = 2`,
}
if testing.Verbose() {
ir.Pretty(os.Stderr, policy)
for _, d := range planner.Debug() {
fmt.Fprint(os.Stderr, d)
}
}
})
}
Expand Down Expand Up @@ -926,10 +934,13 @@ func TestOptimizeLookup(t *testing.T) {
p := New()
p.vars.Put(ast.Var("x"), p.newLocal())

_, _, _, _, opt := p.optimizeLookup(r, ast.MustParseRef("data.foo[x].bar.q"))
if exp, act := false, opt; exp != act {
rulesets, _, _, _, opt := p.optimizeLookup(r, ast.MustParseRef("data.foo[x].bar.q"))
if exp, act := true, opt; exp != act {
t.Errorf("expected 'optimize' %v, got %v\n", exp, act)
}
if exp, act := 0, len(rulesets); exp != act {
t.Fatalf("expected %d rulesets, got %d\n", exp, act)
}
})
}

Expand Down
40 changes: 21 additions & 19 deletions internal/wasm/sdk/opa/opa_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ package opa_test
import (
"context"
"fmt"
"os"
"testing"

"github.com/open-policy-agent/opa/ast"
Expand Down Expand Up @@ -37,8 +38,8 @@ func TestOPA(t *testing.T) {
Policy: `a = true`,
Query: "data.p.a = x",
Evals: []Eval{
Eval{Result: `{{"x": true}}`},
Eval{Result: `{{"x": true}}`},
{Result: `{{"x": true}}`},
{Result: `{{"x": true}}`},
},
WantErr: "",
},
Expand All @@ -47,8 +48,8 @@ func TestOPA(t *testing.T) {
Policy: `a = input`,
Query: "data.p.a = x",
Evals: []Eval{
Eval{Input: "false", Result: `{{"x": false}}`},
Eval{Input: "true", Result: `{{"x": true}}`},
{Input: "false", Result: `{{"x": false}}`},
{Input: "true", Result: `{{"x": true}}`},
},
WantErr: "",
},
Expand All @@ -58,8 +59,8 @@ func TestOPA(t *testing.T) {
Query: "data.p.a = x",
Data: `{"q": false}`,
Evals: []Eval{
Eval{Result: `{{"x": false}}`},
Eval{NewData: `{"q": true}`, Result: `{{"x": true}}`},
{Result: `{{"x": false}}`},
{NewData: `{"q": true}`, Result: `{{"x": true}}`},
},
WantErr: "",
},
Expand All @@ -69,8 +70,8 @@ func TestOPA(t *testing.T) {
Query: "data.p.a = x",
Data: `{"q": false, "r": true}`,
Evals: []Eval{
Eval{Result: `{{"x": false}}`},
Eval{NewPolicy: `a = data.r`, Result: `{{"x": true}}`},
{Result: `{{"x": false}}`},
{NewPolicy: `a = data.r`, Result: `{{"x": true}}`},
},
WantErr: "",
},
Expand All @@ -80,8 +81,8 @@ func TestOPA(t *testing.T) {
Query: "data.p.a = x",
Data: `{"q": 0, "r": 1}`,
Evals: []Eval{
Eval{Result: `{{"x": 0}}`},
Eval{NewPolicy: `a = data.r`, NewData: `{"q": 2, "r": 3}`, Result: `{{"x": 3}}`},
{Result: `{{"x": 0}}`},
{NewPolicy: `a = data.r`, NewData: `{"q": 2, "r": 3}`, Result: `{{"x": 3}}`},
},
WantErr: "",
},
Expand All @@ -90,8 +91,8 @@ func TestOPA(t *testing.T) {
Policy: `a = count(data.q) + sum(data.q)`,
Query: "data.p.a = x",
Evals: []Eval{
Eval{NewData: `{"q": []}`, Result: `{{"x": 0}}`},
Eval{NewData: `{"q": [1, 2]}`, Result: `{{"x": 5}}`},
{NewData: `{"q": []}`, Result: `{{"x": 0}}`},
{NewData: `{"q": [1, 2]}`, Result: `{{"x": 5}}`},
},
WantErr: "",
},
Expand All @@ -100,7 +101,7 @@ func TestOPA(t *testing.T) {
Policy: `a = true`,
Query: "data.p.b = x",
Evals: []Eval{
Eval{Result: `set()`},
{Result: `set()`},
},
WantErr: "",
},
Expand Down Expand Up @@ -166,8 +167,8 @@ a = "c" { input > 2 }`,
}`,
Query: "data.p.hello = x",
Evals: []Eval{
Eval{Input: `{"message": "xxxxxxx"}`, Result: `{{"x": false}}`},
Eval{Input: `{"message": "world"}`, Result: `{{"x": true}}`},
{Input: `{"message": "xxxxxxx"}`, Result: `{{"x": false}}`},
{Input: `{"message": "world"}`, Result: `{{"x": true}}`},
},
},
{
Expand All @@ -179,22 +180,22 @@ a = "c" { input > 2 }`,
}`,
Query: "data.p.hello = x",
Evals: []Eval{
Eval{Input: `{"message": "xxxxxxx"}`, Result: `{{"x": false}}`},
Eval{Input: `{"message": "world"}`, Result: `{{"x": true}}`},
{Input: `{"message": "xxxxxxx"}`, Result: `{{"x": false}}`},
{Input: `{"message": "world"}`, Result: `{{"x": true}}`},
},
},
{
Description: "regex.match with pattern from input",
Query: `x = regex.match(input.re, "foo")`,
Evals: []Eval{
Eval{Input: `{"re": "^foo$"}`, Result: `{{"x": true}}`},
{Input: `{"re": "^foo$"}`, Result: `{{"x": true}}`},
},
},
{
Description: "regex.find_all_string_submatch_n with pattern from input",
Query: `x = regex.find_all_string_submatch_n(input.re, "-axxxbyc-", -1)`,
Evals: []Eval{
Eval{Input: `{"re": "a(x*)b(y|z)c"}`, Result: `{{"x":[["axxxbyc","xxx","y"]]}}`},
{Input: `{"re": "a(x*)b(y|z)c"}`, Result: `{{"x":[["axxxbyc","xxx","y"]]}}`},
},
},
{
Expand Down Expand Up @@ -382,6 +383,7 @@ func compileRegoToWasm(policy string, query string) []byte {
cr, err := rego.New(
rego.Query(query),
rego.Module("module.rego", module),
rego.Dump(os.Stderr),
).Compile(context.Background(), rego.CompilePartial(false))
if err != nil {
panic(err)
Expand Down
11 changes: 8 additions & 3 deletions rego/rego.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1285,7 +1285,7 @@ func (r *Rego) Compile(ctx context.Context, opts ...CompileOption) (*CompileResu

const queryName = "eval" // NOTE(tsandall): the query name is arbitrary

policy, err := planner.New().
p := planner.New().
WithQueries([]planner.QuerySet{
{
Name: queryName,
Expand All @@ -1294,8 +1294,8 @@ func (r *Rego) Compile(ctx context.Context, opts ...CompileOption) (*CompileResu
},
}).
WithModules(modules).
WithBuiltinDecls(decls).
Plan()
WithBuiltinDecls(decls)
policy, err := p.Plan()
if err != nil {
return nil, err
}
Expand All @@ -1305,6 +1305,11 @@ func (r *Rego) Compile(ctx context.Context, opts ...CompileOption) (*CompileResu
fmt.Fprintln(r.dump, "-----")
ir.Pretty(r.dump, policy)
fmt.Fprintln(r.dump)

fmt.Fprintln(r.dump, "planner debug:")
for _, d := range p.Debug() {
fmt.Fprintln(r.dump, d)
}
}

m, err := wasm.New().WithPolicy(policy).Compile()
Expand Down

0 comments on commit dcdbf36

Please sign in to comment.