Implement dependency analysis package for AST elements

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 {

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
+	})
+}

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
+		}
+	}
+}

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