From f5956ba34f2a2814fd86d7082f439833d575674f Mon Sep 17 00:00:00 2001 From: Matthew Mussomele Date: Mon, 3 Jul 2017 12:40:23 -0700 Subject: [PATCH] Implement dependency analysis package for AST elements --- ast/term.go | 11 + dependencies/deps.go | 313 ++++++++++++++++++++++++++++ dependencies/deps_test.go | 416 ++++++++++++++++++++++++++++++++++++++ dependencies/doc.go | 7 + 4 files changed, 747 insertions(+) create mode 100644 dependencies/deps.go create mode 100644 dependencies/deps_test.go create mode 100644 dependencies/doc.go diff --git a/ast/term.go b/ast/term.go index 79cc16cb52..5641e754b7 100644 --- a/ast/term.go +++ b/ast/term.go @@ -742,6 +742,17 @@ func (ref Ref) HasPrefix(other Ref) bool { return true } +// ConstantPrefix returns the constant portion of the ref starting from the head. +func (ref Ref) ConstantPrefix() Ref { + ref = ref.Copy() + + i := ref.Dynamic() + if i < 0 { + return ref + } + return ref[:i] +} + // 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/dependencies/deps.go b/dependencies/deps.go new file mode 100644 index 0000000000..f19f05c113 --- /dev/null +++ b/dependencies/deps.go @@ -0,0 +1,313 @@ +// Copyright 2017 The OPA Authors. All rights reserved. +// Use of this source code is governed by an Apache2 +// license that can be found in the LICENSE file. + +package dependencies + +import ( + "fmt" + "sort" + + "github.com/open-policy-agent/opa/ast" + "github.com/open-policy-agent/opa/util" +) + +// All returns the list of data ast.Refs that the given AST element depends on. +func All(x interface{}) (resolved []ast.Ref, err error) { + var rawResolved []ast.Ref + switch x := x.(type) { + case *ast.Module, *ast.Package, *ast.Import, *ast.Rule, *ast.Head, ast.Body, *ast.Expr, *ast.With, *ast.Term, ast.Ref, ast.Object, ast.Array, *ast.Set, *ast.ArrayComprehension: + default: + return nil, fmt.Errorf("not an ast element: %v", x) + } + + visitor := ast.NewGenericVisitor(func(x interface{}) bool { + switch x := x.(type) { + case *ast.Package, *ast.Import: + return true + case *ast.Module, *ast.Head, *ast.Expr, *ast.With, *ast.Term, ast.Object, ast.Array, *ast.Set, *ast.ArrayComprehension: + case *ast.Rule: + rawResolved = append(rawResolved, ruleDeps(x)...) + return true + case ast.Body: + vars := ast.NewVarVisitor() + ast.Walk(vars, x) + + var arr ast.Array + for v := range vars.Vars() { + if v.IsWildcard() { + continue + } + arr = append(arr, ast.NewTerm(v)) + } + + // The analysis will discard variables that are not used in + // direct comparisions or in the output. Since lone Bodies are + // often queries, we want all the variables to be in the output. + r := &ast.Rule{ + Head: &ast.Head{Name: ast.Var("_"), Value: ast.NewTerm(arr)}, + Body: x, + } + rawResolved = append(rawResolved, ruleDeps(r)...) + return true + case ast.Ref: + rawResolved = append(rawResolved, x) + } + return false + }) + ast.Walk(visitor, x) + if len(rawResolved) == 0 { + return nil, nil + } + + sort.Slice(rawResolved, func(i, j int) bool { + return rawResolved[i].Compare(rawResolved[j]) < 0 + }) + + return filter(rawResolved, func(a, b ast.Ref) bool { + return a.Compare(b) == 0 + }), nil +} + +// Minimal returns the list of data ast.Refs that the given AST element depends on. +// If an AST element depends on a ast.Ref that is a prefix of another dependency, the +// ast.Ref that is the prefix of the other will be the only one in the returned list. +// +// As an example, if an element depends on data.x and data.x.y, only data.x will +// be in the returned list. +func Minimal(x interface{}) (resolved []ast.Ref, err error) { + rawResolved, err := All(x) + if err != nil { + return nil, err + } + + if len(rawResolved) == 0 { + return nil, nil + } + + return filter(rawResolved, func(a, b ast.Ref) bool { + return b.HasPrefix(a) + }), nil +} + +// filter removes all items from the list that cause pref to return true. It is +// called on adjacent pairs of elements, and the one passed as the second argument +// to pref is considered the current one being examined. The first argument will +// be the element immediately preceding it. +func filter(rs []ast.Ref, pred func(ast.Ref, ast.Ref) bool) (filtered []ast.Ref) { + if len(rs) == 0 { + return nil + } + + last := rs[0] + filtered = append(filtered, last) + for i := 1; i < len(rs); i++ { + cur := rs[i] + if pred(last, cur) { + continue + } + + filtered = append(filtered, cur) + last = cur + } + + return filtered +} + +func ruleDeps(rule *ast.Rule) (resolved []ast.Ref) { + vars, others := extractEq(rule.Body) + joined := joinVarRefs(vars) + + headVars := rule.Head.Vars() + headRefs, others := resolveOthers(others, headVars, joined) + + resolveRef := func(r ast.Ref) bool { + resolved = append(resolved, expandRef(r, joined)...) + return false + } + + varVisitor := ast.NewVarVisitor().WithParams(ast.VarVisitorParams{SkipRefHead: true}) + // Clean up whatever refs are remaining among the other expressions. + for _, expr := range others { + ast.WalkRefs(expr, resolveRef) + ast.Walk(varVisitor, expr) + } + + // If a reference ending in a header variable is a prefix of an already + // resolved reference, skip it and simply walk the nodes below it. + visitor := skipVisitor{resolveRef} + for _, r := range headRefs { + if !containsPrefix(resolved, r) { + resolved = append(resolved, r.Copy()) + } + ast.Walk(visitor, r) + } + + usedVars := varVisitor.Vars() + resolveRemainingVars(joined, visitor, usedVars, headVars) + return resolved +} + +// Extract the equality expressions from each rule, they contain +// the potential split references. In order to be considered for +// joining, an equality must have a variable on one side and a +// reference on the other. Any other construct is thrown into +// the others list to be resolved later. +func extractEq(exprs ast.Body) (vars map[ast.Var][]ast.Ref, others []*ast.Expr) { + vars = map[ast.Var][]ast.Ref{} + for v := range exprs.Vars(ast.VarVisitorParams{}) { + vars[v] = nil + } + + for _, expr := range exprs { + if !expr.IsEquality() { + others = append(others, expr) + continue + } + + terms := expr.Terms.([]*ast.Term) + left, right := terms[1], terms[2] + if l, ok := left.Value.(ast.Var); ok { + if r, ok := right.Value.(ast.Ref); ok { + vars[l] = append(vars[l], r) + continue + } + } else if r, ok := right.Value.(ast.Var); ok { + if l, ok := left.Value.(ast.Ref); ok { + vars[r] = append(vars[r], l) + continue + } + } + + others = append(others, expr) + } + return vars, others +} + +func expandRef(r ast.Ref, vars map[ast.Var]*util.HashMap) []ast.Ref { + head, rest := r[0], r[1:] + if ast.DefaultRootDocument.Equal(head) { + return []ast.Ref{r} + } + + h := head.Value.(ast.Var) + rs, ok := vars[h] + if !ok { + return nil + } + + var expanded []ast.Ref + rs.Iter(func(a, _ util.T) bool { + ref := a.(ast.Ref) + expanded = append(expanded, append(ref.Copy(), rest...)) + return false + }) + return expanded +} + +func joinVarRefs(vars map[ast.Var][]ast.Ref) map[ast.Var]*util.HashMap { + joined := map[ast.Var]*util.HashMap{} + for v := range vars { + joined[v] = util.NewHashMap(refEq, refHash) + } + + done := false + for !done { + done = true + for v, rs := range vars { + for _, r := range rs { + head, rest := r[0], r[1:] + if ast.DefaultRootDocument.Equal(head) { + if _, ok := joined[v].Get(r); !ok { + joined[v].Put(r, struct{}{}) + done = false + } + continue + } + + h, ok := head.Value.(ast.Var) + if !ok { + panic("not reached") + } + + joined[h].Iter(func(a, _ util.T) bool { + jr := a.(ast.Ref) + join := append(jr.Copy(), rest...) + if _, ok := joined[v].Get(join); !ok { + joined[v].Put(join, struct{}{}) + done = false + } + return false + }) + } + } + } + + return joined +} + +func resolveOthers(others []*ast.Expr, headVars ast.VarSet, joined map[ast.Var]*util.HashMap) (headRefs []ast.Ref, leftover []*ast.Expr) { + for _, expr := range others { + if term, ok := expr.Terms.(*ast.Term); ok { + if r, ok := term.Value.(ast.Ref); ok { + end := r[len(r)-1] + v, ok := end.Value.(ast.Var) + if ok && headVars.Contains(v) { + headRefs = append(headRefs, expandRef(r, joined)...) + continue + } + } + } + + leftover = append(leftover, expr) + } + + return headRefs, leftover +} + +func resolveRemainingVars(joined map[ast.Var]*util.HashMap, visitor ast.Visitor, usedVars ast.VarSet, headVars ast.VarSet) { + for v, refs := range joined { + visit := visitor + if headVars.Contains(v) || refs.Len() > 1 || usedVars.Contains(v) { + visit = visit.Visit(nil) + } + + refs.Iter(func(a, _ util.T) bool { + r := a.(ast.Ref) + ast.Walk(visit, r) + return false + }) + } +} + +func containsPrefix(refs []ast.Ref, r ast.Ref) bool { + for _, ref := range refs { + if ref.HasPrefix(r) { + return true + } + } + return false +} + +func refEq(a, b util.T) bool { + ar, aok := a.(ast.Ref) + br, bok := b.(ast.Ref) + return aok && bok && ar.Equal(br) +} + +func refHash(a util.T) int { + return a.(ast.Ref).Hash() +} + +type skipVisitor struct { + fn func(ast.Ref) bool +} + +func (sv skipVisitor) Visit(v interface{}) ast.Visitor { + return ast.NewGenericVisitor(func(x interface{}) bool { + if r, ok := x.(ast.Ref); ok { + return sv.fn(r) + } + return false + }) +} diff --git a/dependencies/deps_test.go b/dependencies/deps_test.go new file mode 100644 index 0000000000..d159eaedc3 --- /dev/null +++ b/dependencies/deps_test.go @@ -0,0 +1,416 @@ +// Copyright 2017 The OPA Authors. All rights reserved. +// Use of this source code is governed by an Apache2 +// license that can be found in the LICENSE file. + +package dependencies + +import ( + "fmt" + "sort" + "testing" + + "github.com/open-policy-agent/opa/ast" +) + +type testData struct { + ast string + min []string + full []string +} + +func TestDependencies(t *testing.T) { + tests := []testData{ + { + ast: `package a.b.c + import data.a.x + import data.a.y + + d { + a = x + b = data.a.y + a = b + }`, + min: []string{"a.x", "a.y"}, + }, + { + ast: `package a.b.c + import data.a.x + + d { + a = x + b = a.y + a = "4" + }`, + + min: []string{"a.x"}, + }, + { + ast: `package a.b.c + import data.a.x + + d { + true = a.y + a = x + }`, + + min: []string{"a.x.y"}, + }, + { + ast: `package a.b.c + import data.a.x + + d = f { + a = x.y + e = "foo" + f = [b | a[i].b = e + a[i].c = b] + }`, + + min: []string{"a.x.y[i].b", "a.x.y[i].c"}, + }, + { + ast: `package a.b.c + import data.a.x + + d[i] = f { + x[i] + count([1 | x[i].foo = "foo"], f) + }`, + + min: []string{"a.x[i].foo"}, + }, + { + ast: `package a.b.c + import data.a.x + + d[i] = f { + count([1 | x[i].foo = "foo"], f) + x[i] + }`, + + min: []string{"a.x[i].foo"}, + }, + { + ast: `package a.b.c + import data.a.x + + d[i] = f { + b = x.y + b[i] + count([1 | x.y[_].foo = "foo"], f) + }`, + + min: []string{"a.x.y[i]", "a.x.y[_].foo"}, + }, + { + ast: `package a.b.c + import data.a.x + + d[i] = f { + b = x.y + b[i] + count([1 | x.y[0].foo = "foo"], f) + }`, + + min: []string{"a.x.y[i]", "a.x.y[0].foo"}, + }, + { + ast: `package a.b.c + import data.a.x + + d[i] { + x[i] + }`, + + min: []string{"a.x[i]"}, + }, + { + ast: `package a.b.c + import data.a.x + + d[i] { + b = x.y + b[i] + }`, + + min: []string{"a.x.y[i]"}, + }, + { + ast: `package a.b.c + import data.a.x + + d[a] = b { + a = x.y + b = x.z + x.y.z = "foo" + x.z.a.b = "fizz" + x.a.b = "bar" + }`, + + min: []string{"a.x.y", "a.x.z", "a.x.a.b"}, + full: []string{"a.x.y.z", "a.x.z.a.b"}, + }, + { + ast: `package a.b.c + import data.a.x + + f { + a = x + b = a.y + c = b.z + d = c.a + e = d.b + e.c = "foo" + }`, + + min: []string{"a.x.y.z.a.b.c"}, + }, + { + ast: `package a.b.c + import data.a.x + import data.a.y + import data.j + + f { + a = x + b = a.y + c = b.z + d = c.a + e = d.b + d = j + e.c = "foo" + a = y + e["foo"] = "bar" + }`, + + min: []string{"a.x", "a.y", "j"}, + full: []string{ + "a.x.y", + "a.x.y.z", + "a.x.y.z.a", + "a.x.y.z.a.b", + "a.x.y.z.a.b.c", + "a.x.y.z.a.b.foo", + "a.y.y", + "a.y.y.z", + "a.y.y.z.a", + "a.y.y.z.a.b", + "a.y.y.z.a.b.c", + "a.y.y.z.a.b.foo", + "j.b", + "j.b.c", + "j.b.foo", + }, + }, + { + ast: `package a.b.c + import data.a.x + import data.a.y + import data.j + + f { + a = x + b = a.y + c = b.z + d = c.a + d.b = "foo" + } + + g { + a = x.z.b + e = x + h = e.y + d = h.z + i = d.j + a = 9 + i = "bar" + }`, + + min: []string{ + "a.x.y.z.a.b", + "a.x.y.z.j", + "a.x.z.b", + }, + }, + { + ast: `package a.b.c + import data.a.x + import data.a.y + import data.a.z + import data.a.f + import data.a.g + + a[i] = [j, k] { + b = x.y + y[b.z[0]] = "foo" + z[b.a.c][i] + f[j][b.a.b] = "foo" + g["foo"][k] + }`, + + min: []string{ + "a.x.y.z[0]", + "a.x.y.a.c", + "a.x.y.a.b", + "a.y[b.z[0]]", + "a.z[b.a.c][i]", + "a.f[j][b.a.b]", + "a.g.foo[k]", + }, + }, + { + ast: `package a.b.c + import data.a.x + + f { + a = x + b = a.y + b = a.z + c = b.i + c = b.j + d = c.m + d = c.n + d.f = "foo" + }`, + + min: []string{"a.x.y", "a.x.z"}, + full: []string{ + "a.x.y.i", + "a.x.y.i.m", + "a.x.y.i.m.f", + "a.x.y.i.n", + "a.x.y.i.n.f", + "a.x.y.j", + "a.x.y.j.m", + "a.x.y.j.m.f", + "a.x.y.j.n", + "a.x.y.j.n.f", + "a.x.z.i", + "a.x.z.i.m", + "a.x.z.i.m.f", + "a.x.z.i.n", + "a.x.z.i.n.f", + "a.x.z.j", + "a.x.z.j.m", + "a.x.z.j.m.f", + "a.x.z.j.n", + "a.x.z.j.n.f", + }, + }, + { + ast: `package a.b.c + import data.a.x + + f = [z, g] { + a = x + b = a.y.z + indexof(a.b, b, z) + c = b.e + indexof(c, a.b[a.c].c, g) + }`, + + min: []string{"a.x.y.z", "a.x.b", "a.x.c"}, + full: []string{"a.x.b[a.c].c", "a.x.y.z.e"}, + }, + { + ast: `package a.b.c + import data.a.x + + f = [g] { + a = x + b = a.y.z + c = b.e + d = a.u + indexof(b, a.c[d].c, g) + c = "foo" + }`, + + min: []string{"a.x.y.z", "a.x.u", "a.x.c[d].c"}, + full: []string{"a.x.y.z.e"}, + }, + { + ast: `package a.b.c + import data.a.x + + f { + x + }`, + + min: []string{"a.x"}, + }, + } + + for n, test := range tests { + t.Run(fmt.Sprint(n), func(t *testing.T) { + module := ast.MustParseModule(test.ast) + compiler := ast.NewCompiler() + if compiler.Compile(map[string]*ast.Module{"test": module}); compiler.Failed() { + t.Fatalf("Failed to compile policy: %v", compiler.Errors) + } + + var exp []ast.Ref + for _, e := range test.min { + r := ast.MustParseRef("data." + e) + exp = append(exp, r) + } + sort.Slice(exp, func(i, j int) bool { + return exp[i].Compare(exp[j]) < 0 + }) + + mod := compiler.Modules["test"] + min, full := runDeps(t, mod, test) + + // Test that we get the same result by analyzing all the + // rules separately. + var minRules, fullRules []ast.Ref + for _, rule := range mod.Rules { + m, f := runDeps(t, rule, test) + minRules = append(minRules, m...) + fullRules = append(fullRules, f...) + } + + assertRefSliceEq(t, exp, min) + assertRefSliceEq(t, exp, minRules) + + for _, full := range test.full { + r := ast.MustParseRef("data." + full) + exp = append(exp, r) + } + sort.Slice(exp, func(i, j int) bool { + return exp[i].Compare(exp[j]) < 0 + }) + + assertRefSliceEq(t, exp, full) + assertRefSliceEq(t, exp, fullRules) + }) + } +} + +func runDeps(t *testing.T, x interface{}, test testData) (min, full []ast.Ref) { + min, err := Minimal(x) + if err != nil { + t.Fatalf("Unexpected dependency error: %v", err) + } + + full, err = All(x) + if err != nil { + t.Fatalf("Unexpected dependency error: %v", err) + } + + return min, full +} + +// For some reason, reflect.DeepEqual doesn't work on Ref slices. +func assertRefSliceEq(t *testing.T, exp, result []ast.Ref) { + if len(result) != len(exp) { + t.Fatalf("Expected refs %v, got %v", exp, result) + } + + for i, e := range exp { + r := result[i] + if e.Compare(r) != 0 { + t.Fatalf("Expected refs %v, got %v", exp, result) + break + } + } +} diff --git a/dependencies/doc.go b/dependencies/doc.go new file mode 100644 index 0000000000..a37c4ae3ef --- /dev/null +++ b/dependencies/doc.go @@ -0,0 +1,7 @@ +// Copyright 2017 The OPA Authors. All rights reserved. +// Use of this source code is governed by an Apache2 +// license that can be found in the LICENSE file. + +// Package dependencies provides functions for determining the set of ast.Refs that AST +// elements depend on. +package dependencies