Add safety check on closures/array comprehensions

This change set does away with the old way of determining which variables are
outputs. Equality is now handled with special care. Outputs that would make a
variable safe by depending on another unsafe variable are no longer included.
As a result, the occurs check in the topdown implementation is no longer
needed. This change was introduced to handle odd cases involving
comprehensions, e.g., x = y, x = [ y | y = 1 ]. In this case, without exluding
unsafe vars, the query would evaluate with x/[1]. This would violate the
semantics, because in the comprehension y/1.

Also, fix bug in reordering whereby potentially unsafe expressions were added
to the reordered body multiple times. This occurred because the expression
would be added once when the preceeding expression made it safe and then again
once the outer loop got it. With the fix, we reprocess the body each time an
expression is added to the reordered body.

ast/builtins.go

@@ -29,10 +29,10 @@ var BuiltinMap map[Var]*Builtin
 
 // Equality represents the "=" operator.
 var Equality = &Builtin{
-	Name:         Var("="),
-	Alias:        Var("eq"),
-	NumArgs:      2,
-	RecTargetPos: []int{0, 1},
+	Name:      Var("="),
+	Alias:     Var("eq"),
+	NumArgs:   2,
+	TargetPos: []int{0, 1},
 }
 
 // GreaterThan represents the ">" comparison operator.
@@ -73,11 +73,10 @@ var NotEqual = &Builtin{
 // Builtin represents a built-in function supported by OPA. Every
 // built-in function is uniquely identified by a name.
 type Builtin struct {
-	Name         Var
-	Alias        Var
-	NumArgs      int
-	TargetPos    []int
-	RecTargetPos []int
+	Name      Var
+	Alias     Var
+	NumArgs   int
+	TargetPos []int
 }
 
 // GetPrintableName returns a printable name for the builtin.
@@ -91,26 +90,14 @@ func (b *Builtin) GetPrintableName() string {
 	return b.Name.String()
 }
 
-// Unifies returns true if a term in the given position will unify
-// non-recursively or recursively.
-func (b *Builtin) Unifies(i int) bool {
+// IsTargetPos returns true if a variable in the i-th position will be
+// bound when the expression is evaluated.
+func (b *Builtin) IsTargetPos(i int) bool {
 	for _, x := range b.TargetPos {
 		if x == i {
 			return true
 		}
 	}
-	return b.UnifiesRecursively(i)
-}
-
-// UnifiesRecursively returns true if a term in the given position will
-// unify recursively, i.e., variables embedded inside a collection type
-// will unify.
-func (b *Builtin) UnifiesRecursively(i int) bool {
-	for _, x := range b.RecTargetPos {
-		if x == i {
-			return true
-		}
-	}
 	return false
 }
 

ast/builtins_test.go

@@ -10,26 +10,11 @@ import (
 )
 
 func TestUnifies(t *testing.T) {
-	b := &Builtin{Name: Var("dummy"), NumArgs: 4, RecTargetPos: []int{2, 3}, TargetPos: []int{1}}
-	expected := []int{1, 2, 3}
+	b := &Builtin{Name: Var("dummy"), NumArgs: 4, TargetPos: []int{1, 3}}
+	expected := []int{1, 3}
 	result := []int{}
 	for i := 0; i < 4; i++ {
-		if b.Unifies(i) {
-			result = append(result, i)
-		}
-	}
-	if !reflect.DeepEqual(expected, result) {
-		t.Errorf("Expected %v but got: %v", expected, result)
-	}
-}
-
-func TestUnifiesRecursively(t *testing.T) {
-
-	b := &Builtin{Name: Var("dummy"), NumArgs: 4, RecTargetPos: []int{2, 3}, TargetPos: []int{1}}
-	expected := []int{2, 3}
-	result := []int{}
-	for i := 0; i < 4; i++ {
-		if b.UnifiesRecursively(i) {
+		if b.IsTargetPos(i) {
 			result = append(result, i)
 		}
 	}

ast/compile.go

@@ -204,7 +204,7 @@ func (c *Compiler) checkSafetyHead() {
 	for _, m := range c.Modules {
 		for _, r := range m.Rules {
 			headVars := r.HeadVars()
-			bodyVars := r.Body.Vars()
+			bodyVars := r.Body.Vars(true)
 			for headVar := range headVars {
 				if _, ok := bodyVars[headVar]; !ok {
 					c.err("unsafe variable from head of %v: %v", r.Name, headVar)
@@ -533,6 +533,19 @@ func (vs unsafeVars) Add(e *Expr, v Var) {
 	}
 }
 
+func (vs unsafeVars) Set(e *Expr, s VarSet) {
+	vs[e] = s
+}
+
+func (vs unsafeVars) Update(o unsafeVars) {
+	for k, v := range o {
+		if _, ok := vs[k]; !ok {
+			vs[k] = VarSet{}
+		}
+		vs[k].Update(v)
+	}
+}
+
 func (vs unsafeVars) Vars() VarSet {
 	r := VarSet{}
 	for _, s := range vs {
@@ -604,50 +617,149 @@ func findRulesRec(node *ModuleTreeNode, ref Ref) []*Rule {
 // contains a mapping of expressions to unsafe variables in those expressions.
 func reorderBodyForSafety(globals VarSet, body Body) (Body, unsafeVars) {
 
+	body, unsafe := reorderBodyForClosures(globals, body)
+	if len(unsafe) != 0 {
+		return nil, unsafe
+	}
+
 	reordered := Body{}
-	unsafe := unsafeVars{}
+	safe := VarSet{}
 
 	for _, e := range body {
-		for v := range e.Vars() {
-			if !globals.Contains(v) {
+		for v := range e.Vars(true) {
+			if globals.Contains(v) {
+				safe.Add(v)
+			} else {
 				unsafe.Add(e, v)
 			}
 		}
 	}
 
-	safe := VarSet{}
+	for {
+		n := len(reordered)
 
-	for _, e := range body {
+		for _, e := range body {
+			if reordered.Contains(e) {
+				continue
+			}
 
-		safe.Update(e.OutputVars())
+			safe.Update(e.OutputVars(safe))
 
-		for v := range unsafe[e] {
-			if safe.Contains(v) {
-				delete(unsafe[e], v)
+			for v := range unsafe[e] {
+				if safe.Contains(v) {
+					delete(unsafe[e], v)
+				}
+			}
+
+			if len(unsafe[e]) == 0 {
+				delete(unsafe, e)
+				reordered = append(reordered, e)
 			}
 		}
 
-		if len(unsafe[e]) == 0 {
-			reordered = append(reordered, e)
-			delete(unsafe, e)
+		if len(reordered) == n {
+			break
+		}
+	}
 
-			// Check if other expressions in the body are considered safe
-			// now. If they are considered safe now, they can be added
-			// to the end of the re-ordered body.
-			for _, e := range body {
-				if reordered.Contains(e) {
-					continue
-				}
-				for v := range unsafe[e] {
-					if safe.Contains(v) {
-						delete(unsafe[e], v)
-					}
-				}
-				if len(unsafe[e]) == 0 {
-					reordered = append(reordered, e)
-					delete(unsafe, e)
-				}
+	// Recursively visit closures and perform the safety checks on them.
+	// Update the globals at each expression to include the variables that could
+	// be closed over.
+	g := globals.Copy()
+	for i, e := range reordered {
+		if i > 0 {
+			g.Update(reordered[i-1].Vars(true))
+		}
+		vis := &bodySafetyVisitor{
+			current: e,
+			globals: g,
+			unsafe:  unsafe,
+		}
+		Walk(vis, e)
+	}
+
+	return reordered, unsafe
+}
+
+type bodySafetyVisitor struct {
+	current *Expr
+	globals VarSet
+	unsafe  unsafeVars
+}
+
+func (vis *bodySafetyVisitor) Visit(x interface{}) Visitor {
+	switch x := x.(type) {
+	case *Expr:
+		cpy := *vis
+		cpy.current = x
+		return &cpy
+	case *ArrayComprehension:
+		vis.checkArrayComprehensionSafety(x)
+		return nil
+	}
+	return vis
+}
+
+func (vis *bodySafetyVisitor) checkArrayComprehensionSafety(ac *ArrayComprehension) {
+	// Check term for safety. This is analagous to the rule head safety check.
+	tv := ac.Term.Vars()
+	bv := ac.Body.Vars(true)
+	bv.Update(vis.globals)
+	uv := tv.Diff(bv)
+	for v := range uv {
+		vis.unsafe.Add(vis.current, v)
+	}
+
+	// Check body for safety, reordering as necessary.
+	r, u := reorderBodyForSafety(vis.globals, ac.Body)
+	if len(u) == 0 {
+		ac.Body = r
+	} else {
+		vis.unsafe.Update(u)
+	}
+}
+
+// reorderBodyForClosures returns a copy of the body ordered such that
+// expressions (such as array comprehensions) that close over variables are ordered
+// after other expressions that contain the same variable in an output position.
+func reorderBodyForClosures(globals VarSet, body Body) (Body, unsafeVars) {
+
+	reordered := Body{}
+	unsafe := unsafeVars{}
+
+	for {
+		n := len(reordered)
+
+		for _, e := range body {
+			if reordered.Contains(e) {
+				continue
 			}
+
+			// Collect vars that are contained in closures within this
+			// expression.
+			vs := VarSet{}
+			WalkClosures(e, func(x interface{}) bool {
+				vis := &varVisitor{vars: vs}
+				Walk(vis, x)
+				return true
+			})
+
+			// Compute vars that are closed over from the body but not yet
+			// contained in the output position of an expression in the reordered
+			// body. These vars are considered unsafe.
+			cv := vs.Intersect(body.Vars(true)).Diff(globals)
+			uv := cv.Diff(reordered.OutputVars(globals))
+
+			if len(uv) == 0 {
+				reordered = append(reordered, e)
+				delete(unsafe, e)
+			} else {
+				unsafe.Set(e, uv)
+			}
+		}
+
+		if len(reordered) == n {
+			break
 		}
 	}
 

ast/compile_test.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"reflect"
 	"sort"
+	"strings"
 	"testing"
 )
 
@@ -178,31 +179,111 @@ func TestCompilerCheckSafetyHead(t *testing.T) {
 }
 
 func TestCompilerCheckSafetyBodyReordering(t *testing.T) {
-	c := NewCompiler()
-	c.Modules = getCompilerTestModules()
-	c.Modules["newMod"] = MustParseModule(`
-	package a.b
-	needsReorder = true :- a[i] = x, a = [1,2,3,4]
-	needsReorderNegated = true :- a = [true, false], b = [true, false], not a[i], b[i]
-	`)
-	compileStages(c, "", "checkSafetyBody")
+	tests := []struct {
+		note     string
+		body     string
+		expected interface{}
+	}{
+		// trivial cases
+		{"noop", "x = 1, x != 0", "x = 1, x != 0"},
+		{"var/ref", "a[i] = x, a = [1,2,3,4]", "a = [1,2,3,4], a[i] = x"},
+		{"negation",
+			"a = [true, false], b = [true, false], not a[i], b[i]",
+			"a = [true, false], b = [true, false], b[i], not a[i]"},
+		{"built-in", "x != 0, count([1,2,3], x)", "count([1,2,3], x), x != 0"},
+		{"var/var 1", "x = y, z = 1, y = z", "z = 1, y = z, x = y"},
+		{"var/var 2", "x = y, 1 = z, z = y", "1 = z, z = y, x = y"},
+		{"var/var 3", "x != 0, y = x, y = 1", "y = 1, y = x, x != 0"},
 
-	assertNotFailed(t, c)
-
-	expected1 := MustParseBody(`a = [1,2,3,4], a[i] = x`)
-
-	reordered1 := c.Modules["newMod"].Rules[0].Body
-
-	if !expected1.Equal(reordered1) {
-		t.Errorf("Expected body to be re-ordered and equal to %v but got: %v", expected1, reordered1)
+		// comprehensions
+		{"array compr/var", "x != 0, [y | y = 1] = x", "[y | y = 1] = x, x != 0"},
+		{"array compr/array", "[1] != [x], [y | y = 1] = [x]", "[y | y = 1] = [x], [1] != [x]"},
 	}
 
-	expected2 := MustParseBody(`a = [true, false], b = [true, false], b[i], not a[i]`)
+	for i, tc := range tests {
+		c := NewCompiler()
+		c.Modules = map[string]*Module{
+			"mod": MustParseModule(
+				fmt.Sprintf(`package test
+						 p :- %s`, tc.body)),
+		}
 
-	reordered2 := c.Modules["newMod"].Rules[1].Body
+		compileStages(c, "", "checkSafetyBody")
+		switch exp := tc.expected.(type) {
+		case string:
+			if c.Failed() {
+				t.Errorf("%v (#%d): Unexpected compilation error: %v", tc.note, i, c.FlattenErrors())
+				return
+			}
+			e := MustParseBody(exp)
+			if !e.Equal(c.Modules["mod"].Rules[0].Body) {
+				t.Errorf("%v (#%d): Expected body to be ordered and equal to %v but got: %v", tc.note, i, e, c.Modules["mod"].Rules[0].Body)
+			}
+		case error:
+			if len(c.Errors) > 0 {
+				if !reflect.DeepEqual(c.Errors[0], exp) {
+					t.Errorf("%v (#%d): Expected compiler error %v but got: %v", tc.note, i, exp, c.Errors[0])
+				}
+			} else {
+				t.Errorf("%v (#%d): Expected compiler error but got: %v", tc.note, i, c.Modules["mod"].Rules[0])
+			}
+		}
+	}
+}
 
-	if !expected2.Equal(reordered2) {
-		t.Errorf("Expected body to be re-ordered and equal to %v but got: %v", expected2, reordered2)
+func TestCompilerCheckSafetyBodyReorderingClosures(t *testing.T) {
+	c := NewCompiler()
+	c.Modules = map[string]*Module{
+		"mod": MustParseModule(
+			`
+			package compr
+
+			import data.b
+			import data.c
+
+			p :- v     = [null | true],                               # leave untouched
+				 xs    = [x | a[i] = x, a = [y | y != 1, y = c[j]]],  # close over 'i' and 'j', 2-level reorder
+				 xs[j] > 0,
+				 b[i]  = j
+
+			# test that reordering is not performed when closing over different globals, e.g.,
+			# built-ins, data, imports.
+			q :- _ = [x | x = b[i]],
+				 _ = b[j],
+
+				 _ = [x | x = true, x != false],
+				 true != false,
+
+				 _ = [x | data.foo[_] = x],
+				 data.foo[_] = _
+			`),
+	}
+
+	compileStages(c, "", "checkSafetyBody")
+	assertNotFailed(t, c)
+
+	result1 := c.Modules["mod"].Rules[0].Body
+	expected1 := MustParseBody(`
+		v     = [null | true],
+		b[i]  = j,
+		xs    = [x | a = [y | y = c[j], y != 1], a[i] = x],
+		xs[j] > 0
+	`)
+	if !result1.Equal(expected1) {
+		t.Errorf("Expected reordered body to be equal to:\n%v\nBut got:\n%v", expected1, result1)
+	}
+
+	result2 := c.Modules["mod"].Rules[1].Body
+	expected2 := MustParseBody(`
+		_ = [x | x = b[i]],
+		_ = b[j],
+		_ = [x | x = true, x != false],
+		true != false,
+		_ = [x | data.foo[_] = x],
+		data.foo[_] = _
+	`)
+	if !result2.Equal(expected2) {
+		t.Errorf("Expected pre-ordered body to equal:\n%v\nBut got:\n%v", expected2, result2)
 	}
 }
 
@@ -240,7 +321,10 @@ func TestCompilerCheckSafetyBodyErrors(t *testing.T) {
 	unboundNegated3[x] = true :- a = [1,2,3,4], b = [1,2,3,4], not a[i] = x, not b[j] = x
 
 	# i and j would be unbound even though they are in embedded references
-	unboundNegated4 =  true :- a = [{"foo": ["bar", "baz"]}], not a[0].foo = [a[0].foo[i], a[0].foo[j]]
+	unboundNegated4 = true :- a = [{"foo": ["bar", "baz"]}], not a[0].foo = [a[0].foo[i], a[0].foo[j]]
+
+	# x would be unbound as input to count
+	unsafeBuiltin :- count([1,2,x], x)
 
 	# i and x would be bound in the last expression so the third expression is safe
 	negatedSafe = true :- a = [1,2,3,4], b = [1,2,3,4], not a[i] = x, b[i] = x
@@ -248,6 +332,17 @@ func TestCompilerCheckSafetyBodyErrors(t *testing.T) {
 	# x would be unbound because it does not appear in the target position of any expression
 	unboundNoTarget = true :- x > 0, x <= 3, x != 2
 
+	unboundArrayComprBody1 :- _ = [x | x = data.a[_], y > 1]
+	unboundArrayComprBody2 :- _ = [x | x = a[_], a = [y | y = data.a[_], z > 1]]
+	unboundArrayComprBody3 :- _ = [v | v = [x | x = data.a[_]], x > 1]
+	unboundArrayComprTerm1 :- _ = [u | true]
+	unboundArrayComprTerm2 :- _ = [v | v = [w | w != 0]]
+	unboundArrayComprTerm3 :- _ = [x[i] | x = []]
+	unboundArrayComprMixed1 :- _ = [x | y = [a | a = z[i]]]
+
+	unsafeClosure1 :- x = [x | x = 1]
+	unsafeClosure2 :- x = y, x = [y | y = 1]
+
 	negatedImport1 = true :- not foo
 	negatedImport2 = true :- not bar
 	negatedImport3 = true :- not baz
@@ -262,11 +357,29 @@ func TestCompilerCheckSafetyBodyErrors(t *testing.T) {
 		fmt.Errorf("unsafe variables in unboundNegated2: [i x]"),
 		fmt.Errorf("unsafe variables in unboundNegated3: [i j x]"),
 		fmt.Errorf("unsafe variables in unboundNegated4: [i j]"),
+		fmt.Errorf("unsafe variables in unsafeBuiltin: [x]"),
 		fmt.Errorf("unsafe variables in unboundNoTarget: [x]"),
+		fmt.Errorf("unsafe variables in unboundArrayComprBody1: [y]"),
+		fmt.Errorf("unsafe variables in unboundArrayComprBody2: [z]"),
+		fmt.Errorf("unsafe variables in unboundArrayComprBody3: [x]"),
+		fmt.Errorf("unsafe variables in unboundArrayComprTerm1: [u]"),
+		fmt.Errorf("unsafe variables in unboundArrayComprTerm2: [w]"),
+		fmt.Errorf("unsafe variables in unboundArrayComprTerm3: [i]"),
+		fmt.Errorf("unsafe variables in unboundArrayComprMixed1: [x z]"),
+		fmt.Errorf("unsafe variables in unsafeClosure1: [x]"),
+		fmt.Errorf("unsafe variables in unsafeClosure2: [y]"),
 	}
 
 	if !reflect.DeepEqual(expected, c.Errors) {
-		t.Errorf("Expected %v but got:%v", expected, c.Errors)
+		e := []string{}
+		for _, x := range expected {
+			e = append(e, x.Error())
+		}
+		r := []string{}
+		for _, x := range c.Errors {
+			r = append(r, x.Error())
+		}
+		t.Errorf("Expected:\n%v\nBut got:\n%v", strings.Join(e, "\n"), strings.Join(r, "\n"))
 	}
 
 }

ast/policy.go

@@ -238,6 +238,16 @@ func (body Body) IsGround() bool {
 	return true
 }
 
+// OutputVars returns a VarSet containing the variables that would be bound by evaluating
+// the body.
+func (body Body) OutputVars(safe VarSet) VarSet {
+	o := safe.Copy()
+	for _, e := range body {
+		o.Update(e.OutputVars(o))
+	}
+	return o.Diff(safe)
+}
+
 func (body Body) String() string {
 	var buf []string
 	for _, v := range body {
@@ -246,10 +256,13 @@ func (body Body) String() string {
 	return strings.Join(buf, ", ")
 }
 
-// Vars returns map where keys represent all of the variables found in the
-// body. The values of the map are ignored.
-func (body Body) Vars() VarSet {
-	vis := &varVisitor{vars: VarSet{}}
+// Vars returns a VarSet containing all of the variables in the body. If skipClosures is true,
+// variables contained inside closures within the body will be ignored.
+func (body Body) Vars(skipClosures bool) VarSet {
+	vis := &varVisitor{
+		vars:         VarSet{},
+		skipClosures: skipClosures,
+	}
 	Walk(vis, body)
 	return vis.vars
 }
@@ -327,46 +340,24 @@ func (expr *Expr) IsGround() bool {
 	return true
 }
 
-// OutputVars returns the set of variables that would be bound by
-// evaluating this expression in isolation.
-func (expr *Expr) OutputVars() VarSet {
-
-	result := VarSet{}
-	if expr.Negated {
-		return result
-	}
-
-	vis := &varVisitor{
-		skipRefHead:    true,
-		skipObjectKeys: true,
-		vars:           VarSet{},
-	}
-
-	switch ts := expr.Terms.(type) {
-	case *Term:
-		if r, ok := ts.Value.(Ref); ok {
-			Walk(vis, r)
-		}
-	case []*Term:
-		b := BuiltinMap[ts[0].Value.(Var)]
-		for i, t := range ts[1:] {
-			switch v := t.Value.(type) {
-			case Object, Array:
-				if b.UnifiesRecursively(i) {
-					Walk(vis, v)
+// OutputVars returns a VarSet containing variables that would be bound by evaluating
+// this expression.
+func (expr *Expr) OutputVars(safe VarSet) VarSet {
+	if !expr.Negated {
+		switch terms := expr.Terms.(type) {
+		case *Term:
+			return expr.outputVarsRefs()
+		case []*Term:
+			name := terms[0].Value.(Var)
+			if b := BuiltinMap[name]; b != nil {
+				if b.Name.Equal(Equality.Name) {
+					return expr.outputVarsEquality(safe)
 				}
-			case Var:
-				if b.Unifies(i) {
-					result.Add(v)
-				}
-			case Ref:
-				Walk(vis, v)
+				return expr.outputVarsBuiltins(b, safe)
 			}
 		}
 	}
-
-	result.Update(vis.vars)
-	return result
+	return VarSet{}
 }
 
 func (expr *Expr) String() string {
@@ -404,12 +395,73 @@ func (expr *Expr) UnmarshalJSON(bs []byte) error {
 }
 
 // Vars returns a VarSet containing all of the variables in the expression.
-func (expr *Expr) Vars() VarSet {
-	vis := &varVisitor{vars: VarSet{}}
+// If skipClosures is true then variables contained inside closures within this
+// expression will not be included in the VarSet.
+func (expr *Expr) Vars(skipClosures bool) VarSet {
+	vis := &varVisitor{
+		skipClosures: skipClosures,
+		vars:         VarSet{},
+	}
 	Walk(vis, expr)
 	return vis.vars
 }
 
+func (expr *Expr) outputVarsBuiltins(b *Builtin, safe VarSet) VarSet {
+
+	o := expr.outputVarsRefs()
+	terms := expr.Terms.([]*Term)
+
+	// Check that all input terms are ground or safe.
+	for i, t := range terms[1:] {
+		if b.IsTargetPos(i) {
+			continue
+		}
+		if t.Value.IsGround() {
+			continue
+		}
+		vis := &varVisitor{
+			skipClosures:     true,
+			skipObjectKeys:   true,
+			skipRefHead:      true,
+			skipBuiltinNames: true,
+			vars:             VarSet{},
+		}
+		Walk(vis, t)
+		unsafe := vis.vars.Diff(o).Diff(safe)
+		if len(unsafe) > 0 {
+			return VarSet{}
+		}
+	}
+
+	// Add vars in target positions to result.
+	for i, t := range terms[1:] {
+		if v, ok := t.Value.(Var); ok {
+			if b.IsTargetPos(i) {
+				o.Add(v)
+			}
+		}
+	}
+
+	return o
+}
+
+func (expr *Expr) outputVarsEquality(safe VarSet) VarSet {
+	ts := expr.Terms.([]*Term)
+	o := expr.outputVarsRefs()
+	o.Update(safe)
+	o.Update(Unify(o, ts[1], ts[2]))
+	return o.Diff(safe)
+}
+
+func (expr *Expr) outputVarsRefs() VarSet {
+	o := VarSet{}
+	WalkRefs(expr, func(r Ref) bool {
+		o.Update(r.OutputVars())
+		return false
+	})
+	return o
+}
+
 // NewBuiltinExpr creates a new Expr object with the supplied terms.
 // The builtin operator must be the first term.
 func NewBuiltinExpr(terms ...*Term) *Expr {
@@ -417,9 +469,11 @@ func NewBuiltinExpr(terms ...*Term) *Expr {
 }
 
 type varVisitor struct {
-	skipRefHead    bool
-	skipObjectKeys bool
-	vars           VarSet
+	skipRefHead      bool
+	skipObjectKeys   bool
+	skipClosures     bool
+	skipBuiltinNames bool
+	vars             VarSet
 }
 
 func (vis *varVisitor) Visit(v interface{}) Visitor {
@@ -439,6 +493,22 @@ func (vis *varVisitor) Visit(v interface{}) Visitor {
 			return nil
 		}
 	}
+	if vis.skipClosures {
+		switch v.(type) {
+		case *ArrayComprehension:
+			return nil
+		}
+	}
+	if vis.skipBuiltinNames {
+		if v, ok := v.(*Expr); ok {
+			if ts, ok := v.Terms.([]*Term); ok {
+				for _, t := range ts[1:] {
+					Walk(vis, t)
+				}
+				return nil
+			}
+		}
+	}
 	if v, ok := v.(Var); ok {
 		vis.vars.Add(v)
 	}

ast/policy_test.go

@@ -140,12 +140,39 @@ func TestBodyIsGround(t *testing.T) {
 }
 
 func TestExprOutputVars(t *testing.T) {
-	body := MustParseBody(`{"a": [{x: y}, b[z]]} = c[i], [{"a": d[j][k]}] != xs`)
-	one := body[0]
-	vars := one.OutputVars()
-	expected := NewVarSet(Var("y"), Var("z"), Var("i"))
-	if !reflect.DeepEqual(expected, vars) {
-		t.Errorf("Expected output vars %v from %v but got: %v", expected, one, vars)
+
+	tests := []struct {
+		note     string
+		expr     string
+		safe     string
+		expected string
+	}{
+		{"ref 1", "a[i].b[j]", "[a]", "[i, j]"},
+		{"ref 2", "[1,2,a[i]]", "[a]", "[i]"},
+		{"simple unify", `{"a": [{x: y}, b[z]]} = c[i]`, "[b, c]", "[y, z, i]"},
+		{"built-in", "count([], x)", "[]", "[x]"},
+	}
+
+	for i, tc := range tests {
+
+		expr := MustParseBody(tc.expr)[0]
+		safe := VarSet{}
+		for _, x := range MustParseTerm(tc.safe).Value.(Array) {
+			safe.Add(x.Value.(Var))
+		}
+
+		result := expr.OutputVars(safe)
+
+		expected := VarSet{}
+		for _, x := range MustParseTerm(tc.expected).Value.(Array) {
+			expected.Add(x.Value.(Var))
+		}
+
+		missing := expected.Diff(result)
+		extra := result.Diff(expected)
+		if len(missing) != 0 || len(extra) != 0 {
+			t.Errorf("%s (%d): Missing output vars: %v, extra output vars: %v", tc.note, i, missing, extra)
+		}
 	}
 }
 

ast/term.go

@@ -132,6 +132,13 @@ func (term *Term) UnmarshalJSON(bs []byte) error {
 	return nil
 }
 
+// Vars returns a VarSet with variables contained in this term.
+func (term *Term) Vars() VarSet {
+	vis := &varVisitor{vars: VarSet{}}
+	Walk(vis, term)
+	return vis.vars
+}
+
 // Null represents the null value defined by JSON.
 type Null struct{}
 
@@ -415,6 +422,17 @@ func (ref Ref) Underlying() ([]interface{}, error) {
 	return r, nil
 }
 
+// OutputVars returns a VarSet containing variables that would be bound by evaluating
+//  this expression in isolation.
+func (ref Ref) OutputVars() VarSet {
+	vis := &varVisitor{
+		vars:        VarSet{},
+		skipRefHead: true,
+	}
+	Walk(vis, ref)
+	return vis.vars
+}
+
 // QueryIterator defines the interface for querying AST documents with references.
 type QueryIterator func(map[Var]Value, Value) error
 

ast/unify.go

@@ -0,0 +1,165 @@
+// Copyright 2016 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 ast
+
+// Unify returns a set of variables that will be unified when the equality expression defined by
+// terms a and b is evaluated. The unifier assumes that variables in the VarSet safe are already
+// unified.
+func Unify(safe VarSet, a *Term, b *Term) VarSet {
+	u := &unifier{
+		safe:    safe,
+		unified: VarSet{},
+		unknown: map[Var]VarSet{},
+	}
+	u.unify(a, b)
+	return u.unified
+}
+
+type unifier struct {
+	safe    VarSet
+	unified VarSet
+	unknown map[Var]VarSet
+}
+
+func (u *unifier) isSafe(x Var) bool {
+	return u.safe.Contains(x) || u.unified.Contains(x)
+}
+
+func (u *unifier) unify(a *Term, b *Term) {
+
+	switch a := a.Value.(type) {
+
+	case Var:
+		switch b := b.Value.(type) {
+		case Var:
+			if u.isSafe(b) {
+				u.markSafe(a)
+			} else if u.isSafe(a) {
+				u.markSafe(b)
+			} else {
+				u.markUnknown(a, b)
+				u.markUnknown(b, a)
+			}
+		case Array, Object:
+			u.unifyAll(a, b)
+		default:
+			u.markSafe(a)
+		}
+
+	case Ref:
+		switch b := b.Value.(type) {
+		case Var:
+			u.markSafe(b)
+		case Array, Object:
+			u.markAllSafe(b, a)
+		}
+
+	case *ArrayComprehension:
+		switch b := b.Value.(type) {
+		case Var:
+			u.markSafe(b)
+		case Array:
+			u.markAllSafe(b, a)
+		}
+
+	case Array:
+		switch b := b.Value.(type) {
+		case Var:
+			u.unifyAll(b, a)
+		case Ref, *ArrayComprehension:
+			u.markAllSafe(a, b)
+		case Array:
+			if len(a) == len(b) {
+				for i := range a {
+					u.unify(a[i], b[i])
+				}
+			}
+		}
+
+	case Object:
+		switch b := b.Value.(type) {
+		case Var:
+			u.unifyAll(b, a)
+		case Ref:
+			u.markAllSafe(a, b)
+		case Object:
+			if len(a) == len(b) {
+				for i := range a {
+					u.unify(a[i][1], b[i][1])
+				}
+			}
+		}
+
+	default:
+		switch b := b.Value.(type) {
+		case Var:
+			u.markSafe(b)
+		}
+	}
+}
+
+func (u *unifier) markAllSafe(x Value, y Value) {
+	vis := u.varVisitor()
+	Walk(vis, x)
+	for v := range vis.vars {
+		u.markSafe(v)
+	}
+}
+
+func (u *unifier) markSafe(x Var) {
+	u.unified.Add(x)
+
+	// Add dependencies of 'x' to safe set
+	vs := u.unknown[x]
+	delete(u.unknown, x)
+	for v := range vs {
+		u.markSafe(v)
+	}
+
+	// Add dependants of 'x' to safe set if they have no more
+	// dependencies.
+	for v, deps := range u.unknown {
+		if deps.Contains(x) {
+			delete(deps, x)
+			if len(deps) == 0 {
+				u.markSafe(v)
+			}
+		}
+	}
+}
+
+func (u *unifier) markUnknown(a, b Var) {
+	if _, ok := u.unknown[a]; !ok {
+		u.unknown[a] = NewVarSet()
+	}
+	u.unknown[a].Add(b)
+}
+
+func (u *unifier) unifyAll(a Var, b Value) {
+	if u.isSafe(a) {
+		u.markAllSafe(b, a)
+	} else {
+		vis := u.varVisitor()
+		Walk(vis, b)
+		unsafe := vis.vars.Diff(u.safe).Diff(u.unified)
+		if len(unsafe) == 0 {
+			u.markSafe(a)
+		} else {
+			for v := range unsafe {
+				u.markUnknown(a, v)
+			}
+		}
+	}
+}
+
+func (u *unifier) varVisitor() *varVisitor {
+	return &varVisitor{
+		skipRefHead:      true,
+		skipObjectKeys:   true,
+		skipClosures:     true,
+		skipBuiltinNames: true,
+		vars:             VarSet{},
+	}
+}

ast/unify_test.go

@@ -0,0 +1,78 @@
+// Copyright 2016 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 ast
+
+import "testing"
+
+func TestUnify(t *testing.T) {
+
+	tests := []struct {
+		note     string
+		expr     string
+		safe     string
+		expected string
+	}{
+		// collection cases
+		{"array/ref", "[1,2,x] = a[_]", "[a]", "[x]"},
+		{"array/ref (reversed)", "a[_] = [1,2,x]", "[a]", "[x]"},
+		{"array/var", "[1,2,x] = y", "[x]", "[y]"},
+		{"array/var (reversed)", "y = [1,2,x]", "[x]", "[y]"},
+		{"array/var-2", "[1,2,x] = y", "[y]", "[x]"},
+		{"array/var-2 (reversed)", "y = [1,2,x]", "[y]", "[x]"},
+		{"array/uneven", "[1,2,x] = [y,x]", "[]", "[]"},
+		{"array/uneven-2", "[1,2,x] = [y,x]", "[x]", "[]"},
+		{"object/ref", `{"x": x} = a[_]`, "[a]", "[x]"},
+		{"object/ref (reversed)", `a[_] = {"x": x}`, "[a]", "[x]"},
+		{"object/var", `{"x": 1, "y": x} = y`, "[x]", "[y]"},
+		{"object/var (reversed)", `y = {"x": 1, "y": x}`, "[x]", "[y]"},
+		{"object/var-2", `{"x": 1, "y": x} = y`, "[y]", "[x]"},
+		{"object/var-3", `{"x": 1, "y": x} = y`, "[]", "[]"},
+		{"object/uneven", `{"x": x, "y": 1} = {"x": y}`, "[]", "[]"},
+		{"object/uneven", `{"x": x, "y": 1} = {"x": y}`, "[x]", "[]"},
+
+		// transitive cases
+		{"trans/redundant", "[x, x] = [x, 0]", "[]", "[x]"},
+		{"trans/simple", "[x, 1] = [y, y]", "[]", "[y, x]"},
+		{"trans/array", "[x, y] = [y, [z, a]]", "[x]", "[a, y, z]"},
+		{"trans/object", `[x, y] = [y, {"a":a,"z":z}]`, "[x]", "[a, y, z]"},
+		{"trans/ref", "[x, y, [x, y, i]] = [1, a[i], z]", "[a, i]", "[x, y, z]"},
+		{"trans/lazy", "[x, z, 2] = [1, [y, x], y]", "[]", "[x, y, z]"},
+		{"trans/redundant-nested", "[x, z, z] = [1, [y, x], [2, 1]]", "[]", "[x, y, z]"},
+		{"trans/bidirectional", "[x, z, y] = [[z,y], [1,y], 2]", "[]", "[x, y, z]"},
+		{"trans/occurs", "[x, z, y] = [[y,z], [y, 1], [2, x]]", "[]", "[]"},
+	}
+
+	for i, tc := range tests {
+
+		expr := MustParseBody(tc.expr)[0]
+		safe := VarSet{}
+		for _, x := range MustParseTerm(tc.safe).Value.(Array) {
+			safe.Add(x.Value.(Var))
+		}
+
+		terms := expr.Terms.([]*Term)
+		if !terms[0].Value.Equal(Equality.Name) {
+			panic(terms)
+		}
+
+		a, b := terms[1], terms[2]
+		unified := Unify(safe, a, b)
+		result := VarSet{}
+		for k := range unified {
+			result.Add(k)
+		}
+
+		expected := VarSet{}
+		for _, x := range MustParseTerm(tc.expected).Value.(Array) {
+			expected.Add(x.Value.(Var))
+		}
+
+		missing := expected.Diff(result)
+		extra := result.Diff(expected)
+		if len(missing) != 0 || len(extra) != 0 {
+			t.Errorf("%s (%d): Missing vars: %v, extra vars: %v", tc.note, i, missing, extra)
+		}
+	}
+}

ast/varset.go

@@ -41,6 +41,28 @@ func (s VarSet) Copy() VarSet {
 	return cpy
 }
 
+// Diff returns a VarSet containing variables in s that are not in vs.
+func (s VarSet) Diff(vs VarSet) VarSet {
+	r := VarSet{}
+	for v := range s {
+		if !vs.Contains(v) {
+			r.Add(v)
+		}
+	}
+	return r
+}
+
+// Intersect returns a VarSet containing variables in s that are in vs.
+func (s VarSet) Intersect(vs VarSet) VarSet {
+	r := VarSet{}
+	for v := range s {
+		if vs.Contains(v) {
+			r.Add(v)
+		}
+	}
+	return r
+}
+
 // Update merges the other VarSet into this VarSet.
 func (s VarSet) Update(vs VarSet) {
 	for v := range vs {

ast/visit.go

@@ -77,3 +77,43 @@ func Walk(v Visitor, x interface{}) {
 		Walk(w, x.Body)
 	}
 }
+
+// WalkClosures calls the function f on all closures under x. If the function f
+// returns true, AST nodes under the last node will not be visited.
+func WalkClosures(x interface{}, f func(interface{}) bool) {
+	vis := &GenericVisitor{func(x interface{}) bool {
+		switch x.(type) {
+		case *ArrayComprehension:
+			return f(x)
+		}
+		return false
+	}}
+	Walk(vis, x)
+}
+
+// WalkRefs calls the function f on all references under x. If the function f
+// returns true, AST nodes under the last node will not be visited.
+func WalkRefs(x interface{}, f func(Ref) bool) {
+	vis := &GenericVisitor{func(x interface{}) bool {
+		if r, ok := x.(Ref); ok {
+			return f(r)
+		}
+		return false
+	}}
+	Walk(vis, x)
+}
+
+// GenericVisitor implements the Visitor interface to provide
+// a utility to walk over AST nodes using a closure. If the closure
+// returns true, the visitor will not walk over AST nodes under x.
+type GenericVisitor struct {
+	f func(x interface{}) bool
+}
+
+// Visit calls the function f on the GenericVisitor.
+func (vis *GenericVisitor) Visit(x interface{}) Visitor {
+	if vis.f(x) {
+		return nil
+	}
+	return vis
+}

topdown/topdown.go

@@ -78,15 +78,6 @@ func (ctx *Context) BindVar(variable ast.Var, value ast.Value) *Context {
 	if variable.Equal(value) {
 		return ctx
 	}
-	occurs := walkValue(value, func(other ast.Value) bool {
-		if variable.Equal(other) {
-			return true
-		}
-		return false
-	})
-	if occurs {
-		return nil
-	}
 
 	cpy := *ctx
 	cpy.Locals = storage.NewBindings()
@@ -389,23 +380,6 @@ func dereferenceVar(v ast.Var, ctx *Context) (interface{}, error) {
 func evalContext(ctx *Context, iter Iterator) error {
 
 	if ctx.Index >= len(ctx.Query) {
-
-		// Check if the bindings contain values that are non-ground. E.g.,
-		// suppose the query's final expression is "x = y" and "x" and "y"
-		// do not appear elsewhere in the query. In this case, "x" and "y"
-		// will be bound to each other; they will not be ground and so
-		// the proof should not be considered successful.
-		isNonGround := ctx.Locals.Iter(func(k, v ast.Value) bool {
-			if !v.IsGround() {
-				return true
-			}
-			return false
-		})
-
-		if isNonGround {
-			return nil
-		}
-
 		ctx.traceFinish()
 		return iter(ctx)
 	}

topdown/topdown_test.go

@@ -234,7 +234,6 @@ func TestTopDownCompleteDoc(t *testing.T) {
 		{`object/nested composites: {"a": [1], "b": [2], "c": [3]}`,
 			`p = {"a": [1], "b": [2], "c": [3]} :- true`,
 			`{"a": [1], "b": [2], "c": [3]}`},
-		{"var/var", "p = true :- x = y", ""},
 	}
 
 	data := loadSmallTestData()
@@ -258,7 +257,6 @@ func TestTopDownPartialSetDoc(t *testing.T) {
 		{"nested composites", "p[x] :- f[i] = x", `[{"xs": [1.0], "ys": [2.0]}, {"xs": [2.0], "ys": [3.0]}]`},
 		{"deep ref/heterogeneous", "p[x] :- c[i][j][k] = x", `[null, 3.14159, true, false, true, false, "foo"]`},
 		{"composite var value", "p[x] :- x = [i, a[i]]", "[[0,1],[1,2],[2,3],[3,4]]"},
-		{"var/var", "p[x] :- x = y", "[]"},
 	}
 
 	data := loadSmallTestData()
@@ -278,8 +276,6 @@ func TestTopDownPartialObjectDoc(t *testing.T) {
 		{"composites", "p[k] = v :- d[k] = v", `{"e": ["bar", "baz"]}`},
 		{"non-string key", "p[k] = v :- a[k] = v", fmt.Errorf("illegal object key type float64: 0")},
 		{"body/join var", "p[k] = v :- a[i] = v, g[k][i] = v", `{"a": 1, "b": 2, "c": 4}`},
-		{"var/var key", "p[k] = v :- v = 1, k = x", "{}"},
-		{"var/var val", `p[k] = v :- k = "x", v = x`, "{}"},
 	}
 
 	data := loadSmallTestData()
@@ -310,8 +306,6 @@ func TestTopDownEqExpr(t *testing.T) {
 		{"undefined: array deep var 2", "p = true :- [[1,x],[3,4]] = [[1,2],[x,4]]", ""},
 		{"undefined: array uneven", `p = true :- [true, false, "foo", "deadbeef"] = c[i][j]`, ""},
 		{"undefined: object uneven", `p = true :- {"a": 1, "b": 2} = {"a": 1}`, ""},
-		{"undefined: occurs 1", "p = true :- [y,x] = [[x],y]", ""},
-		{"undefined: occurs 2", "p = true :- [y,x] = [{\"a\": x}, y]", ""},
 
 		// ground terms
 		{"ground: bool", `p = true :- true = true`, "true"},