From c9d1a8db1ffecf1d85dd6f9992daca3b77605a35 Mon Sep 17 00:00:00 2001 From: Johan Fylling Date: Wed, 27 Sep 2023 14:06:19 +0200 Subject: [PATCH] planner: Adding support for general ref rule heads (#6235) Fixes: #5995 Signed-off-by: Johan Fylling --- ast/term.go | 14 ++ internal/planner/planner.go | 131 +++++++++++++++--- internal/planner/planner_test.go | 41 ++++++ internal/planner/rules.go | 36 ++++- internal/wasm/sdk/test/e2e/exceptions.yaml | 12 +- internal/wasm/sdk/test/e2e/external_test.go | 4 + .../test-partialobjectdoc-ref.yaml | 22 +++ .../testdata/refheads/test-generic-refs.yaml | 117 +++++++++++++++- .../refheads/test-refs-as-rule-heads.yaml | 12 ++ 9 files changed, 353 insertions(+), 36 deletions(-) create mode 100644 test/cases/testdata/partialobjectdoc/test-partialobjectdoc-ref.yaml diff --git a/ast/term.go b/ast/term.go index 7c50b14877..9ef8ae7474 100644 --- a/ast/term.go +++ b/ast/term.go @@ -1032,6 +1032,20 @@ func (ref Ref) ConstantPrefix() Ref { return ref[:i] } +func (ref Ref) StringPrefix() Ref { + r := ref.Copy() + + for i := 1; i < len(ref); i++ { + switch r[i].Value.(type) { + case String: // pass + default: // cut off + return r[:i] + } + } + + return r +} + // GroundPrefix returns the ground portion of the ref starting from the head. By // definition, the head of the reference is always ground. func (ref Ref) GroundPrefix() Ref { diff --git a/internal/planner/planner.go b/internal/planner/planner.go index 9c3fd7332d..61addb2efd 100644 --- a/internal/planner/planner.go +++ b/internal/planner/planner.go @@ -25,6 +25,8 @@ type QuerySet struct { } type planiter func() error +type planLocalIter func(ir.Local) error +type stmtFactory func(ir.Local) ir.Stmt // Planner implements a query planner for Rego queries. type Planner struct { @@ -147,32 +149,31 @@ func (p *Planner) buildFunctrie() error { } for _, rule := range module.Rules { - r := rule.Ref() - switch r[len(r)-1].Value.(type) { - case ast.String: // pass - default: // cut off - r = r[:len(r)-1] - } + r := rule.Ref().StringPrefix() val := p.rules.LookupOrInsert(r) + + val.rules = val.DescendantRules() val.rules = append(val.rules, rule) + val.children = nil } } return nil } func (p *Planner) planRules(rules []*ast.Rule) (string, error) { + // We know the rules with closer to the root (shorter static path) are ordered first. pathRef := rules[0].Ref() // figure out what our rules' collective name/path is: // if we're planning both p.q.r and p.q[s], we'll name // the function p.q (for the mapping table) - // TODO(sr): this has to change when allowing `p[v].q.r[w]` ref rules - // including the mapping lookup structure and lookup functions pieces := len(pathRef) for i := range rules { r := rules[i].Ref() - if _, ok := r[len(r)-1].Value.(ast.String); !ok { - pieces = len(r) - 1 + for j, t := range r { + if _, ok := t.Value.(ast.String); !ok && j > 0 && j < pieces { + pieces = j + } } } // control if p.a = 1 is to return 1 directly; or insert 1 under key "a" into an object @@ -236,7 +237,11 @@ func (p *Planner) planRules(rules []*ast.Rule) (string, error) { fn.Blocks = append(fn.Blocks, p.blockWithStmt(&ir.MakeObjectStmt{Target: fn.Return})) } case ast.MultiValue: - fn.Blocks = append(fn.Blocks, p.blockWithStmt(&ir.MakeSetStmt{Target: fn.Return})) + if buildObject { + fn.Blocks = append(fn.Blocks, p.blockWithStmt(&ir.MakeObjectStmt{Target: fn.Return})) + } else { + fn.Blocks = append(fn.Blocks, p.blockWithStmt(&ir.MakeSetStmt{Target: fn.Return})) + } } // For complete document rules, allocate one local variable for output @@ -252,6 +257,12 @@ func (p *Planner) planRules(rules []*ast.Rule) (string, error) { var defaultRule *ast.Rule var ruleLoc *location.Location + // We sort rules by ref length, to ensure that when merged, we can detect conflicts when one + // rule attempts to override values (deep and shallow) defined by another rule. + sort.Slice(rules, func(i, j int) bool { + return len(rules[i].Ref()) > len(rules[j].Ref()) + }) + // Generate function blocks for rules. for i := range rules { @@ -320,18 +331,19 @@ func (p *Planner) planRules(rules []*ast.Rule) (string, error) { switch rule.Head.RuleKind() { case ast.SingleValue: if buildObject { - ref := rule.Head.Ref() - last := ref[len(ref)-1] - return p.planTerm(last, func() error { - key := p.ltarget - return p.planTerm(rule.Head.Value, func() error { - value := p.ltarget - p.appendStmt(&ir.ObjectInsertOnceStmt{ - Object: fn.Return, - Key: key, - Value: value, + ref := rule.Ref() + return p.planTerm(rule.Head.Value, func() error { + value := p.ltarget + return p.planNestedObjects(fn.Return, ref[pieces:len(ref)-1], func(obj ir.Local) error { + return p.planTerm(ref[len(ref)-1], func() error { + key := p.ltarget + p.appendStmt(&ir.ObjectInsertOnceStmt{ + Object: obj, + Key: key, + Value: value, + }) + return nil }) - return nil }) }) } @@ -343,6 +355,28 @@ func (p *Planner) planRules(rules []*ast.Rule) (string, error) { return nil }) case ast.MultiValue: + if buildObject { + ref := rule.Ref() + // we drop the trailing set key from the ref + return p.planNestedObjects(fn.Return, ref[pieces:len(ref)-1], func(obj ir.Local) error { + // Last term on rule ref is the key an which the set is assigned in the deepest nested object + return p.planTerm(ref[len(ref)-1], func() error { + key := p.ltarget + return p.planTerm(rule.Head.Key, func() error { + value := p.ltarget + factory := func(v ir.Local) ir.Stmt { return &ir.MakeSetStmt{Target: v} } + return p.planDotOr(obj, key, factory, func(set ir.Local) error { + p.appendStmt(&ir.SetAddStmt{ + Set: set, + Value: value, + }) + p.appendStmt(&ir.ObjectInsertStmt{Key: key, Value: op(set), Object: obj}) + return nil + }) + }) + }) + }) + } return p.planTerm(rule.Head.Key, func() error { p.appendStmt(&ir.SetAddStmt{ Set: fn.Return, @@ -422,6 +456,59 @@ func (p *Planner) planRules(rules []*ast.Rule) (string, error) { return fn.Name, nil } +func (p *Planner) planDotOr(obj ir.Local, key ir.Operand, or stmtFactory, iter planLocalIter) error { + // We're constructing the following plan: + // + // | block a + // | | block b + // | | | dot &{Source:Local Key:{Value:Local} Target:Local} + // | | | break 1 + // | | or &{Target:Local} + // | | *ir.ObjectInsertOnceStmt &{Key:{Value:Local} Value:{Value:Local} Object:Local} + + prev := p.curr + dotBlock := &ir.Block{} + p.curr = dotBlock + + val := p.newLocal() + p.appendStmt(&ir.DotStmt{ + Source: op(obj), + Key: key, + Target: val, + }) + p.appendStmt(&ir.BreakStmt{Index: 1}) + + outerBlock := &ir.Block{ + Stmts: []ir.Stmt{ + &ir.BlockStmt{Blocks: []*ir.Block{dotBlock}}, // FIXME: Set Location + or(val), + &ir.ObjectInsertOnceStmt{Key: key, Value: op(val), Object: obj}, + }, + } + + p.curr = prev + p.appendStmt(&ir.BlockStmt{Blocks: []*ir.Block{outerBlock}}) + return iter(val) +} + +func (p *Planner) planNestedObjects(obj ir.Local, ref ast.Ref, iter planLocalIter) error { + if len(ref) == 0 { + //return fmt.Errorf("nested object construction didn't create object") + return iter(obj) + } + + t := ref[0] + + return p.planTerm(t, func() error { + key := p.ltarget + + factory := func(v ir.Local) ir.Stmt { return &ir.MakeObjectStmt{Target: v} } + return p.planDotOr(obj, key, factory, func(childObj ir.Local) error { + return p.planNestedObjects(childObj, ref[1:], iter) + }) + }) +} + func (p *Planner) planFuncParams(params []ir.Local, args ast.Args, idx int, iter planiter) error { if idx >= len(args) { return iter() diff --git a/internal/planner/planner_test.go b/internal/planner/planner_test.go index cb22d67448..eb6252ee3d 100644 --- a/internal/planner/planner_test.go +++ b/internal/planner/planner_test.go @@ -16,6 +16,7 @@ import ( ) func TestPlannerHelloWorld(t *testing.T) { + t.Setenv("EXPERIMENTAL_GENERAL_RULE_REFS", "true") // NOTE(tsandall): These tests are not meant to give comprehensive coverage // of the planner. Currently we have a suite of end-to-end tests in the @@ -151,6 +152,46 @@ func TestPlannerHelloWorld(t *testing.T) { p[v] = 2 { v := "b" } `}, }, + { + note: "partial object (ref-head) with var", + queries: []string{`data.test.p.q = x`}, + modules: []string{` + package test + p.q.r["a"] = 1 + p.q[v] = 2 { v := "b" } + `}, + }, + { + note: "partial object (ref-head) with var (shallow query)", + queries: []string{`data.test.p = x`}, + modules: []string{` + package test + p.q["a"] = 1 + p.q[v] = 2 { v := "b" } + p.r["c"] = 3 + p.r[v] = 4 { v := "d" } + `}, + }, + { + note: "partial object (ref-head) with var (multiple)", + queries: []string{`data.test.p.q = x`}, + modules: []string{` + package test + p.q["a"] = 1 + p.q[v] = x { l1 := ["b", "c", "d"]; l2 := ["foo", "bar"]; l3 := [2, 3]; v := l1[_]; x := l2[_]; z := l3[_] } + `}, + }, + { + note: "partial object (general ref-head) with var", + queries: []string{`data.test.p.q = x`}, + modules: []string{` + package test + p.q["a"] = 1 + p.q.b.s.baz = 2 + p.q.b.s.foo.c = 3 + p.q[r].s[t].u = v { x := ["foo", "bar"]; r := "b"; t := x[v]} + `}, + }, { note: "every", queries: []string{`data.test.p`}, diff --git a/internal/planner/rules.go b/internal/planner/rules.go index 06fa6d158d..f5d6f3fc6c 100644 --- a/internal/planner/rules.go +++ b/internal/planner/rules.go @@ -98,6 +98,7 @@ func (t *ruletrie) Rules() []*ast.Rule { // // and we're retrieving a.b, we want Rules() to include the rule body // of a.b.c. + // FIXME: We need to go deeper than just immediate children (?) for _, rs := range t.children { if r := rs[len(rs)-1].rules; r != nil { rules = append(rules, r...) @@ -157,13 +158,46 @@ func (t *ruletrie) Lookup(key ast.Ref) *ruletrie { return node } +func (t *ruletrie) LookupShallowest(key ast.Ref) *ruletrie { + node := t + for _, elem := range key { + node = node.Get(elem.Value) + if node == nil { + return nil + } + if len(node.rules) > 0 { + return node + } + } + return node +} + +// TODO: Collapse rules with overlapping extent to same node(?) func (t *ruletrie) LookupOrInsert(key ast.Ref) *ruletrie { - if val := t.Lookup(key); val != nil { + if val := t.LookupShallowest(key); val != nil { + return val } return t.Insert(key) } +func (t *ruletrie) DescendantRules() []*ast.Rule { + if len(t.children) == 0 { + return t.rules + } + + rules := make([]*ast.Rule, len(t.rules), len(t.rules)+len(t.children)) // could be too little + copy(rules, t.rules) + + for _, cs := range t.children { + for _, c := range cs { + rules = append(rules, c.DescendantRules()...) + } + } + + return rules +} + func (t *ruletrie) ChildrenCount() int { return len(t.children) } diff --git a/internal/wasm/sdk/test/e2e/exceptions.yaml b/internal/wasm/sdk/test/e2e/exceptions.yaml index 951a438bd4..664e29b20f 100644 --- a/internal/wasm/sdk/test/e2e/exceptions.yaml +++ b/internal/wasm/sdk/test/e2e/exceptions.yaml @@ -2,14 +2,4 @@ "data/toplevel integer": "https://github.com/open-policy-agent/opa/issues/3711" "data/nested integer": "https://github.com/open-policy-agent/opa/issues/3711" "withkeyword/function: indirect call, arity 1, replacement is value that needs eval (array comprehension)": "https://github.com/open-policy-agent/opa/issues/5311" -"withkeyword/builtin: indirect call, arity 1, replacement is value that needs eval (array comprehension)": "https://github.com/open-policy-agent/opa/issues/5311" -"refheads/general, single var": "Tests with arbitrary vars in rule refs (general refs) are not supported by the planner yet" -"refheads/general, multiple vars": "Tests with arbitrary vars in rule refs (general refs) are not supported by the planner yet" -"refheads/general, deep query": "Tests with arbitrary vars in rule refs (general refs) are not supported by the planner yet" -"refheads/general, overlapping rule, no conflict": "Tests with arbitrary vars in rule refs (general refs) are not supported by the planner yet" -"refheads/general, overlapping rule, conflict": "Tests with arbitrary vars in rule refs (general refs) are not supported by the planner yet" -"refheads/general, set leaf": "Tests with arbitrary vars in rule refs (general refs) are not supported by the planner yet" -"refheads/general, set leaf, deep query": "Tests with arbitrary vars in rule refs (general refs) are not supported by the planner yet" -"refheads/general, input var": "Tests with arbitrary vars in rule refs (general refs) are not supported by the planner yet" -"refheads/general, external non-ground var": "Tests with arbitrary vars in rule refs (general refs) are not supported by the planner yet" -"refheads/general, multiple result-set entries": "Tests with arbitrary vars in rule refs (general refs) are not supported by the planner yet" \ No newline at end of file +"withkeyword/builtin: indirect call, arity 1, replacement is value that needs eval (array comprehension)": "https://github.com/open-policy-agent/opa/issues/5311" \ No newline at end of file diff --git a/internal/wasm/sdk/test/e2e/external_test.go b/internal/wasm/sdk/test/e2e/external_test.go index bdd98526f0..487ee5ae70 100644 --- a/internal/wasm/sdk/test/e2e/external_test.go +++ b/internal/wasm/sdk/test/e2e/external_test.go @@ -65,6 +65,10 @@ func TestWasmE2E(t *testing.T) { t.SkipNow() } + for k, v := range tc.Env { + t.Setenv(k, v) + } + opts := []func(*rego.Rego){ rego.Query(tc.Query), } diff --git a/test/cases/testdata/partialobjectdoc/test-partialobjectdoc-ref.yaml b/test/cases/testdata/partialobjectdoc/test-partialobjectdoc-ref.yaml new file mode 100644 index 0000000000..0fee305bbf --- /dev/null +++ b/test/cases/testdata/partialobjectdoc/test-partialobjectdoc-ref.yaml @@ -0,0 +1,22 @@ +cases: +- note: partialobjectdoc/ref + modules: + - | + package generated + + p.q[k] = v { + k := ["foo", "bar"][v] + } + + p.baz := 2 + + q { + x := "bar" + y := "q" + p[y][x] == 1 + } + + query: data.generated.q = x + want_result: + - x: + true diff --git a/test/cases/testdata/refheads/test-generic-refs.yaml b/test/cases/testdata/refheads/test-generic-refs.yaml index 74fb221f6e..bdbdf9392d 100644 --- a/test/cases/testdata/refheads/test-generic-refs.yaml +++ b/test/cases/testdata/refheads/test-generic-refs.yaml @@ -63,6 +63,41 @@ cases: r: 1 c: r: 2 + - note: 'refheads/general, overlapping rule, different dynamic depths, no conflict' + env: + EXPERIMENTAL_GENERAL_RULE_REFS: "true" + modules: + - | + package test + + p[q].r.s := i { q := ["a", "b", "c"][i] } + p.d.r := 3 + p.e := 4 + query: data.test.p = x + want_result: + - x: + a: + r: + s: 0 + b: + r: + s: 1 + c: + r: + s: 2 + d: + r: 3 + e: 4 + - note: 'refheads/general, self-conflict' + env: + EXPERIMENTAL_GENERAL_RULE_REFS: "true" + modules: + - | + package test + + p[q].r.s := i { q := ["a", "b", "c", "b"][i] } + query: data.test.p = x + want_error_code: eval_conflict_error - note: 'refheads/general, overlapping rule, conflict' env: EXPERIMENTAL_GENERAL_RULE_REFS: "true" @@ -70,10 +105,88 @@ cases: - | package test - p[q].r := i { q := ["a", "b", "c"][i] } + p[q].r.s := i { q := ["a", "b", "c"][i] } p.a.r := 42 query: data.test.p = x - want_error: eval_conflict_error + want_error_code: eval_conflict_error + - note: 'refheads/general, overlapping rule, deep override inside other rule object value, conflict' + env: + EXPERIMENTAL_GENERAL_RULE_REFS: "true" + modules: + - | + package test + + p[q].r := v { + q := "q"; + v := { + "s": { + "t" : 1 + } + } + } + + p.q.r.s.t := 42 + query: data.test.p = x + want_error_code: eval_conflict_error + - note: 'refheads/general, overlapping rule, deep injection into other rule object value, conflict' + env: + EXPERIMENTAL_GENERAL_RULE_REFS: "true" + modules: + - | + package test + + p[q].r := v { + q := "q"; + v := { + "s": { + "t" : 1 + } + } + } + + p.q.r.s.u := 42 + query: data.test.p = x + want_error_code: eval_conflict_error + - note: 'refheads/general, set leaf (shallow ref)' + env: + EXPERIMENTAL_GENERAL_RULE_REFS: "true" + modules: + - | + package test + import future.keywords + + p[q] contains r { + x := ["a", "b", "c"] + q := x[_] + r := x[_] + q != r + } + + p.b contains "foo" + query: data.test.p = x + want_result: + - x: + a: [ "b", "c" ] + b: [ "a", "c", "foo" ] + c: [ "a", "b" ] + - note: 'refheads/general, set leaf (other rule defines dynamic ref portion)' + env: + EXPERIMENTAL_GENERAL_RULE_REFS: "true" + modules: + - | + package test + import future.keywords + + p.q contains r { + r := ["a", "b", "c"][_] + } + + p[q] := r { q := "foo"; r := "bar" } + query: data.test.p = x + want_result: + - x: + q: [ "a", "b", "c" ] + foo: bar - note: 'refheads/general, set leaf' env: EXPERIMENTAL_GENERAL_RULE_REFS: "true" diff --git a/test/cases/testdata/refheads/test-refs-as-rule-heads.yaml b/test/cases/testdata/refheads/test-refs-as-rule-heads.yaml index f8d00b5a7e..8daabc43a0 100644 --- a/test/cases/testdata/refheads/test-refs-as-rule-heads.yaml +++ b/test/cases/testdata/refheads/test-refs-as-rule-heads.yaml @@ -25,6 +25,18 @@ cases: q: r: 1 s: 2 +- note: 'refheads/single-value, with var, conflict' + env: + EXPERIMENTAL_GENERAL_RULE_REFS: "true" + modules: + - | + package test + + p.q.r = 1 + p.q[s] = 2 { s := "r" } + query: data.test.p.q = x + want_error_code: eval_conflict_error + want_error: object keys must be unique - modules: - | package test