Skip to content

Commit

Permalink
wasm: Add support for full virtual document model
Browse files Browse the repository at this point in the history 
Previously the planner only supported references into values generated
by rules or ground references to packages. If the reference was
non-ground the planner would error and references to cached data were
never evaluated.

This commit updates the planner to support the full virtual document
model. References can iterate over packages, merge base and virtual
documents, etc.

Fixes open-policy-agent#1117
Fixes open-policy-agent#1119

Signed-off-by: Torin Sandall <torinsandall@gmail.com>
  • Loading branch information
@tsandall
tsandall committed Oct 11, 2019
1 parent 0abe8dc commit ae66c1a
Show file tree
Hide file tree
Showing 4 changed files with 469 additions and 132 deletions.
308 changes: 233 additions & 75 deletions internal/planner/planner.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1038,7 +1038,9 @@ func (p *Planner) planRef(ref ast.Ref, iter planiter) error {
}

if head.Compare(ast.DefaultRootDocument.Value) == 0 {
return p.planRefData(p.funcs, ref, 0, iter)
virtual := p.funcs.children[ref[0].Value]
base := &baseptr{local: p.vars.GetOrEmpty(ast.DefaultRootDocument.Value.(ast.Var))}
return p.planRefData(virtual, base, ref, 1, iter)
}

p.ltarget, ok = p.vars.Get(head)
Expand Down Expand Up @@ -1078,111 +1080,267 @@ func (p *Planner) planRefRec(ref ast.Ref, index int, iter planiter) error {
})
}

func (p *Planner) planRefData(node *functrie, ref ast.Ref, idx int, iter planiter) error {
type baseptr struct {
local ir.Local
path ast.Ref
}

// planRefData implements the virtual document model by generating the value of
// the ref parameter and invoking the iterator with the planner target set to
// the virtual document and all variables in the reference assigned.
func (p *Planner) planRefData(virtual *functrie, base *baseptr, ref ast.Ref, index int, iter planiter) error {

if idx >= len(ref) {
return p.planRefDataVirtualExtent(node, iter)
// Early-exit if the end of the reference has been reached. In this case the
// plan has to materialize the full extent of the referenced value.
if index >= len(ref) {
return p.planRefDataExtent(virtual, base, iter)
}

term := ref[idx]
// If the reference operand is ground then either continue to the next
// operand or invoke the function for the rule referred to by this operand.
if ref[index].IsGround() {

var vchild *functrie

if virtual != nil {
vchild = virtual.children[ref[index].Value]
}

if vchild != nil && vchild.val != nil {
p.ltarget = p.newLocal()
p.appendStmt(&ir.CallStmt{
Func: vchild.val.Rules[0].Path().String(),
Args: []ir.Local{
p.vars.GetOrEmpty(ast.InputRootDocument.Value.(ast.Var)),
p.vars.GetOrEmpty(ast.DefaultRootDocument.Value.(ast.Var)),
},
Result: p.ltarget,
})
return p.planRefRec(ref, index+1, iter)
}

bchild := *base
bchild.path = append(bchild.path, ref[index])

if _, ok := term.Value.(ast.String); !ok && idx > 0 {
return fmt.Errorf("not implemented: refs with non-string operands")
return p.planRefData(vchild, &bchild, ref, index+1, iter)
}

child, ok := node.children[term.Value]
if !ok {
return nil
exclude := ast.NewSet()

// The planner does not support dynamic dispatch so generate blocks to
// evaluate each of the rulesets on the child nodes.
if virtual != nil {

stmt := &ir.BlockStmt{}

for _, child := range virtual.Children() {

block := &ir.Block{}
prev := p.curr
p.curr = block
key := ast.NewTerm(child)
exclude.Add(key)

// Assignments in each block due to local unification must be undone
// so create a new frame that will be popped after this key is
// processed.
p.vars.Push(map[ast.Var]ir.Local{})

if err := p.planTerm(key, func() error {
return p.planUnifyLocal(p.ltarget, ref[index], func() error {
// Create a copy of the reference with this operand plugged.
// This will result in evaluation of the rulesets on the
// child node.
cpy := ref.Copy()
cpy[index] = key
return p.planRefData(virtual, base, cpy, index, iter)
})
}); err != nil {
return err
}

p.vars.Pop()
p.curr = prev
stmt.Blocks = append(stmt.Blocks, block)
}

p.appendStmt(stmt)
}

if child.val == nil {
return p.planRefData(child, ref, idx+1, iter)
// If the virtual tree was enumerated then we do not want to enumerate base
// trees that are rooted at the same key as any of the virtual sub trees. To
// prevent this we build a set of keys that are to be excluded and check
// below during the base scan.
var lexclude *ir.Local

if exclude.Len() > 0 {
if err := p.planSet(exclude, func() error {
v := p.ltarget
lexclude = &v
return nil
}); err != nil {
return err
}
}

p.ltarget = p.newLocal()
p.ltarget = base.local

// 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 {
return p.planScan(ref[index], func(lkey ir.Local) error {
if lexclude != nil {
lignore := p.newLocal()
p.appendStmt(&ir.NotStmt{
Block: &ir.Block{
Stmts: []ir.Stmt{
&ir.DotStmt{
Source: *lexclude,
Key: lkey,
Target: lignore,
},
},
},
})
}

p.appendStmt(&ir.CallStmt{
Func: ref[:idx+1].String(),
Args: []ir.Local{
p.vars.GetOrEmpty(ast.InputRootDocument.Value.(ast.Var)),
p.vars.GetOrEmpty(ast.DefaultRootDocument.Value.(ast.Var)),
},
Result: p.ltarget,
// Assume that virtual sub trees have been visited already so
// recurse without the virtual node.
return p.planRefData(nil, &baseptr{local: p.ltarget}, ref, index+1, iter)
})
})

return p.planRefRec(ref, idx+1, iter)
}

func (p *Planner) planRefDataVirtualExtent(node *functrie, iter planiter) error {
// planRefDataExtent generates the full extent (combined) of the base and
// virtual nodes and then invokes the iterator with the planner target set to
// the full extent.
func (p *Planner) planRefDataExtent(virtual *functrie, base *baseptr, iter planiter) error {

// Create a new object document. The target is not set until the planner
// recurses so that we can build the hierarchy depth-first.
target := p.newLocal()
vtarget := p.newLocal()

p.appendStmt(&ir.MakeObjectStmt{
Target: target,
})
// Generate the virtual document out of rules contained under the virtual
// node (recursively). This document will _ONLY_ contain values generated by
// rules. No base document values will be included.
if virtual != nil {

for key, child := range node.children {
p.appendStmt(&ir.MakeObjectStmt{
Target: vtarget,
})

// Skip functions.
if child.val != nil && child.val.Arity() > 0 {
continue
}
for key, child := range virtual.children {

lkey := p.newLocal()
idx := p.appendStringConst(string(key.(ast.String)))
p.appendStmt(&ir.MakeStringStmt{
Index: idx,
Target: lkey,
})
// Skip functions.
if child.val != nil && child.val.Arity() > 0 {
continue
}

// Build object hierarchy depth-first.
if child.val == nil {
err := p.planRefDataVirtualExtent(child, func() error {
p.appendStmt(&ir.ObjectInsertStmt{
Object: target,
Key: lkey,
Value: p.ltarget,
})
return nil
lkey := p.newLocal()
idx := p.appendStringConst(string(key.(ast.String)))
p.appendStmt(&ir.MakeStringStmt{
Index: idx,
Target: lkey,
})
if err != nil {
return err

// Build object hierarchy depth-first.
if child.val == nil {
err := p.planRefDataExtent(child, nil, func() error {
p.appendStmt(&ir.ObjectInsertStmt{
Object: vtarget,
Key: lkey,
Value: p.ltarget,
})
return nil
})
if err != nil {
return err
}
continue
}
continue
}

// Generate virtual document for leaf.
lvalue := p.newLocal()

// Add leaf to object if defined.
p.appendStmt(&ir.BlockStmt{
Blocks: []*ir.Block{
&ir.Block{
Stmts: []ir.Stmt{
&ir.CallStmt{
Func: child.val.Rules[0].Path().String(),
Args: []ir.Local{
p.vars.GetOrEmpty(ast.InputRootDocument.Value.(ast.Var)),
p.vars.GetOrEmpty(ast.DefaultRootDocument.Value.(ast.Var)),
// Generate virtual document for leaf.
lvalue := p.newLocal()

// Add leaf to object if defined.
p.appendStmt(&ir.BlockStmt{
Blocks: []*ir.Block{
&ir.Block{
Stmts: []ir.Stmt{
&ir.CallStmt{
Func: child.val.Rules[0].Path().String(),
Args: []ir.Local{
p.vars.GetOrEmpty(ast.InputRootDocument.Value.(ast.Var)),
p.vars.GetOrEmpty(ast.DefaultRootDocument.Value.(ast.Var)),
},
Result: lvalue,
},
&ir.ObjectInsertStmt{
Object: vtarget,
Key: lkey,
Value: lvalue,
},
Result: lvalue,
},
&ir.ObjectInsertStmt{
Object: target,
Key: lkey,
Value: lvalue,
},
},
},
},
})
})
}

// At this point vtarget refers to the full extent of the virtual
// document at ref. If the base pointer is unset, no further processing
// is required.
if base == nil {
p.ltarget = vtarget
return iter()
}
}

// Obtain the base document value and merge (recursively) with the virtual
// document value above if needed.
prev := p.curr
p.curr = &ir.Block{}
p.ltarget = base.local
target := p.newLocal()

err := p.planRefRec(base.path, 0, func() error {

if virtual == nil {
target = p.ltarget
} else {
stmt := &ir.ObjectMergeStmt{
A: p.ltarget,
B: vtarget,
Target: target,
}
p.appendStmt(stmt)
p.appendStmt(&ir.BreakStmt{Index: 1})
}

return nil
})

if err != nil {
return err
}

inner := p.curr
p.curr = &ir.Block{}
p.appendStmt(&ir.BlockStmt{Blocks: []*ir.Block{inner}})

if virtual != nil {
p.appendStmt(&ir.AssignVarStmt{
Source: vtarget,
Target: target,
})
}

// Set target to object and recurse.
outer := p.curr
p.curr = prev
p.appendStmt(&ir.BlockStmt{Blocks: []*ir.Block{outer}})

// At this point, target refers to either the full extent of the base and
// virtual documents at ref or just the base document at ref.
p.ltarget = target

return iter()
}

Expand Down
19 changes: 19 additions & 0 deletions internal/planner/planner_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,25 @@ func TestPlannerHelloWorld(t *testing.T) {
note: "closure",
queries: []string{`a = [1]; {x | a[_] = x}`},
},
{
note: "iteration: packages and rules",
queries: []string{"data.test[x][y] = 3"},
modules: []string{
`
package test.a
p = 1
q = 2 { false }
r = 3
`,
`
package test.z
s = 3
t = 4
`,
},
},
}

for _, tc := range tests {
Expand Down
Loading

0 comments on commit ae66c1a

Please sign in to comment.