Skip to content

Commit

Permalink
Updating virtual partial rule eval
Browse files Browse the repository at this point in the history 
Signed-off-by: Johan Fylling <johan.dev@fylling.se>
  • Loading branch information
@johanfylling
johanfylling committed May 29, 2023
1 parent ddedf1b commit 8ef1c96
Show file tree
Hide file tree
Showing 4 changed files with 241 additions and 0 deletions.
8 changes: 8 additions & 0 deletions ast/index.go
View file
Expand Up @@ -38,6 +38,7 @@ type IndexResult struct {
Default *Rule
EarlyExit bool
OnlyGroundRefs bool
IsDynamic bool
}

// NewIndexResult returns a new IndexResult object.
Expand All @@ -60,6 +61,7 @@ type baseDocEqIndex struct {
defaultRule *Rule
kind RuleKind
onlyGroundRefs bool
isDynamic bool
}

func newBaseDocEqIndex(isVirtual func(Ref) bool) *baseDocEqIndex {
Expand All @@ -68,6 +70,7 @@ func newBaseDocEqIndex(isVirtual func(Ref) bool) *baseDocEqIndex {
isVirtual: isVirtual,
root: newTrieNodeImpl(),
onlyGroundRefs: true,
isDynamic: false,
}
}

Expand All @@ -89,6 +92,9 @@ func (i *baseDocEqIndex) Build(rules []*Rule) bool {
if i.onlyGroundRefs {
i.onlyGroundRefs = rule.Head.Reference.IsGround()
}
if !i.isDynamic {
i.isDynamic = rule.Head.IsDynamic()
}
var skip bool
for _, expr := range rule.Body {
if op := expr.OperatorTerm(); op != nil && i.skipIndexing.Contains(op) {
Expand Down Expand Up @@ -141,6 +147,7 @@ func (i *baseDocEqIndex) Lookup(resolver ValueResolver) (*IndexResult, error) {
result := NewIndexResult(i.kind)
result.Default = i.defaultRule
result.OnlyGroundRefs = i.onlyGroundRefs
result.IsDynamic = i.isDynamic
result.Rules = make([]*Rule, 0, len(tr.ordering))

for _, pos := range tr.ordering {
Expand Down Expand Up @@ -173,6 +180,7 @@ func (i *baseDocEqIndex) AllRules(resolver ValueResolver) (*IndexResult, error)

result := NewIndexResult(i.kind)
result.Default = i.defaultRule
result.IsDynamic = i.isDynamic
result.OnlyGroundRefs = i.onlyGroundRefs
result.Rules = make([]*Rule, 0, len(tr.ordering))

Expand Down
8 changes: 8 additions & 0 deletions ast/policy.go
View file
Expand Up @@ -748,6 +748,7 @@ func (rule *Rule) elseString() string {
return strings.Join(buf, " ")
}


// NewHead returns a new Head object. If args are provided, the first will be
// used for the key and the second will be used for the value.
func NewHead(name Var, args ...*Term) *Head {
Expand Down Expand Up @@ -975,6 +976,13 @@ func (head *Head) SetLoc(loc *Location) {
head.Location = loc
}

// TODO: Find better name for functions with ref-vars in other possitions than the last
func (head *Head) IsDynamic() bool {
pos := head.Reference.Dynamic()
// Ref is dynamic if it has one non-constant term that isn't the first or last term.
return pos > 0 && pos < len(head.Reference) - 1
}

// Copy returns a deep copy of a.
func (a Args) Copy() Args {
cpy := Args{}
Expand Down
60 changes: 60 additions & 0 deletions topdown/eval.go
View file
Expand Up @@ -2386,6 +2386,17 @@ func (e evalVirtualPartial) evalEachRule(iter unifyIterator, unknown bool) error
}

result := e.empty
if e.ir.IsDynamic {
for _, rule := range e.ir.Rules {
result, err = e.evalOneDynamicRefRulePreUnify(rule, result, unknown)
if err != nil {
return err
}
}
e.e.virtualCache.Put(hint.key, result)
return e.evalTerm(iter, e.pos+1, result, e.bindings)
}

for _, rule := range e.ir.Rules {
if err := e.evalOneRulePreUnify(iter, rule, hint, result, unknown); err != nil {
return err
Expand Down Expand Up @@ -2510,6 +2521,55 @@ func (e evalVirtualPartial) evalOneRulePreUnify(iter unifyIterator, rule *ast.Ru
return nil
}

func (e evalVirtualPartial) evalOneDynamicRefRulePreUnify(rule *ast.Rule, result *ast.Term, unknown bool) (*ast.Term, error) {

child := e.e.child(rule.Body)

child.traceEnter(rule)
var defined bool

// Walk the dynamic portion of rule ref to unify vars
err := child.biunifyDynamicRef(e.pos+1, rule.Ref(), e.ref, child.bindings, e.bindings, func() error {
defined = true
return child.eval(func(child *eval) error {
child.traceExit(rule)
var dup bool
var err error
result, dup, err = e.reduce(rule, child.bindings, result)
if err != nil {
return err
} else if !unknown && dup {
child.traceDuplicate(rule)
return nil
}

child.traceRedo(rule)
return nil
})
})

if err != nil {
return nil, err
}

// We're tracing here to exhibit similar behaviour to evalOneRulePreUnify
if !defined {
child.traceFail(rule)
}

return result, nil
}

func (e *eval) biunifyDynamicRef(pos int, a, b ast.Ref, b1, b2 *bindings, iter unifyIterator) error {
if pos >= len(a) || pos >= len(b) {
return iter()
}

return e.biunify(a[pos], b[pos], b1, b2, func() error {
return e.biunifyDynamicRef(pos+1, a, b, b1, b2, iter)
})
}

func (e evalVirtualPartial) evalOneRulePostUnify(iter unifyIterator, rule *ast.Rule) error {

key := e.ref[e.pos+1]
Expand Down
165 changes: 165 additions & 0 deletions topdown/eval_test.go
View file
Expand Up @@ -632,6 +632,7 @@ func TestPartialRule(t *testing.T) {
query: `data = x`,
expErr: "eval_conflict_error: object keys must be unique",
},
// TODO: Add test case with else block
// Overlapping rules
{
note: "partial object with overlapping rule (defining key/value in object)",
Expand Down Expand Up @@ -793,6 +794,170 @@ func TestPartialRule(t *testing.T) {
query: `data = x`,
expErr: "eval_conflict_error: object keys must be unique",
},
// Deep queries
{
note: "deep query into partial object",
module: `package test
p.q[r] := 1 { r := "foo" }
`,
query: `data.test.p.q.foo = x`,
exp: `[{"x": 1}]`,
},
{
note: "deep query to, but not into partial set",
module: `package test
import future.keywords
p.q.r contains s { {"foo", "bar", "bax"}[s] }
`,
query: `data.test.p = x`,
exp: `[{"x": {"q": {"r": ["bar", "bax", "foo"]}}}]`,
},
{
note: "deep query to, but not into partial set",
module: `package test
import future.keywords
p.q.r contains s { {"foo", "bar", "bax"}[s] }
`,
query: `data.test.p.q = x`,
exp: `[{"x": {"r": ["bar", "bax", "foo"]}}]`,
},
{
note: "deep query to, but not into partial set",
module: `package test
import future.keywords
p.q.r contains s { {"foo", "bar", "bax"}[s] }
`,
query: `data.test.p.q.r = x`,
exp: `[{"x": ["bar", "bax", "foo"]}]`,
},
{
note: "deep query into partial set",
module: `package test
import future.keywords
p.q contains r { {"foo", "bar", "bax"}[r] }
`,
query: `data.test.p.q.foo = x`,
exp: `[{"x": "foo"}]`,
},
{
note: "deep query into partial object and object value",
module: `package test
p.q[r] := x {
r := "foo"
x := {"bar": {"baz": 1}}
}
`,
query: `data.test.p.q.foo.bar = x`,
exp: `[{"x": {"baz": 1}}]`,
},
{
note: "deep query into partial object and set value",
module: `package test
import future.keywords
p.q[r].s contains x {
r := "foo"
{"foo", "bar", "bax"}[x]
}
`,
query: `data.test.p.q.foo.s.bar = x`,
exp: `[{"x": "bar"}]`,
},
{
note: "deep query into partial object and object value, non-tail var",
module: `package test
p.q[r].s := x {
r := "foo"
x := {"bar": {"baz": 1}}
}
`,
query: `data.test.p.q.foo.s.bar = x`,
exp: `[{"x": {"baz": 1}}]`,
},
{
note: "deep query into partial object, on first var in ref",
module: `package test
p.q[r].s := 1 { r := "foo" }
`,
query: `data.test.p.q.foo = x`,
exp: `[{"x": {"s": 1}}]`,
},
{
note: "deep query into partial object, beyond first var in ref",
module: `package test
p.q[r].s := 1 { r := "foo" }
`,
query: `data.test.p.q.foo.s = x`,
exp: `[{"x": 1}]`,
},
{
note: "deep query into partial object, on first var in ref, multiple vars",
module: `package test
p.q[r][s] := 1 { r := "foo"; s := "bar" }
`,
query: `data.test.p.q.foo = x`,
exp: `[{"x": {"bar": 1}}]`,
},
{
note: "deep query into partial object, beyond first var in ref, multiple vars",
module: `package test
p.q[r][s] := 1 { r := "foo"; s := "bar" }
`,
query: `data.test.p.q.foo.bar = x`,
exp: `[{"x": 1}]`,
},
{
note: "deep query into partial object, beyond first var in ref, multiple vars",
module: `package test
p.q[r][s].t := 1 { r := "foo"; s := "bar" }
`,
query: `data.test.p.q.foo.bar = x`,
exp: `[{"x": {"t": 1}}]`,
},
{
note: "deep query to partial object, overlapping rules, no dynamic ref",
module: `package test
p.q[r] := 1 { r := "foo" }
p.q.r := 2
`,
query: `data.test.p.q = x`,
exp: `[{"x": {"foo": 1, "r": 2}}]`,
},
{
note: "deep query into partial object, overlapping rules, no dynamic ref",
module: `package test
p.q[r] := 1 { r := "foo" }
p.q.r := 2
`,
query: `data.test.p.q.r = x`,
exp: `[{"x": 2}]`,
},
{
note: "deep query into partial object, overlapping rules, no dynamic ref",
module: `package test
p.q[r] := 1 { r := "foo" }
p.q[r] := 2 { r := "bar" }
`,
query: `data.test.p.q.foo = x`,
exp: `[{"x": 1}]`,
},
{
note: "deep query into partial object, overlapping rules with same key/value, no dynamic ref",
module: `package test
p.q[r] := 1 { r := "foo" }
p.q[r] := 1 { r := "foo" }
`,
query: `data.test.p.q.foo = x`,
exp: `[{"x": 1}]`,
},
{
note: "deep query into partial object, overlapping rules, dynamic ref",
module: `package test
p.q[r].s := 1 { r := "r" }
p.q.r[s] := 2 { s := "foo" }
`,
query: `data.test.p.q.r = x`,
exp: `[{"x": {"s": 1, "foo": 2}}]`,
},
}

for _, tc := range tests {
Expand Down

0 comments on commit 8ef1c96

Please sign in to comment.