diff --git a/ast/compile.go b/ast/compile.go
index 52f4eb8d59..bc9b6bf834 100644
--- a/ast/compile.go
+++ b/ast/compile.go
@@ -62,10 +62,10 @@ type Compiler struct {
 	//                +--- q (1 rule)
 	RuleTree *RuleTreeNode
 
-	// RuleGraph represents the rule dependencies.
-	// An edge (u, v) is added to the graph if rule "u" depends on rule "v".
-	// A rule depends on another rule if it refers to it.
-	RuleGraph map[*Rule]map[*Rule]struct{}
+	// RuleGraph represents the dependencies between rules. An edge (u,v) is
+	// added to the graph if rule "u" depends on rule "v". A rule "u" depends on
+	// rule "v" if rule "u" refers to the virtual document defined by rule "v".
+	RuleGraph *RuleGraph
 
 	generatedVars map[*Module]VarSet
 	moduleLoader  ModuleLoader
@@ -157,7 +157,6 @@ func NewCompiler() *Compiler {
 
 	c := &Compiler{
 		Modules:       map[string]*Module{},
-		RuleGraph:     map[*Rule]map[*Rule]struct{}{},
 		generatedVars: map[*Module]VarSet{},
 	}
 
@@ -370,19 +369,20 @@ func (c *Compiler) checkRecursion() {
 	eq := func(a, b util.T) bool {
 		return a.(*Rule) == b.(*Rule)
 	}
-	for r := range c.RuleGraph {
-		t := &ruleGraphTraveral{
-			graph:   c.RuleGraph,
-			visited: map[*Rule]struct{}{},
-		}
-		if p := util.DFSPath(t, eq, r, r); len(p) > 0 {
-			n := []string{}
-			for _, x := range p {
-				n = append(n, string(x.(*Rule).Head.Name))
+
+	c.RuleTree.DepthFirst(func(node *RuleTreeNode) bool {
+		for _, rule := range node.Rules {
+			t := newRuleGraphTraversal(c.RuleGraph)
+			if p := util.DFSPath(t, eq, rule, rule); len(p) > 0 {
+				n := []string{}
+				for _, x := range p {
+					n = append(n, string(x.(*Rule).Head.Name))
+				}
+				c.err(NewError(RecursionErr, rule.Loc(), "%v %v is recursive: %v", RuleTypeName, rule.Head.Name, strings.Join(n, " -> ")))
 			}
-			c.err(NewError(RecursionErr, r.Loc(), "%v %v is recursive: %v", RuleTypeName, r.Head.Name, strings.Join(n, " -> ")))
 		}
-	}
+		return false
+	})
 }
 
 // checkRuleConflicts ensures that rules definitions are not in conflict.
@@ -656,20 +656,7 @@ func (c *Compiler) setRuleTree() {
 }
 
 func (c *Compiler) setRuleGraph() {
-	for _, m := range c.Modules {
-		for _, r := range m.Rules {
-			edges, ok := c.RuleGraph[r]
-			if !ok {
-				edges = map[*Rule]struct{}{}
-				c.RuleGraph[r] = edges
-			}
-			vis := &ruleGraphBuilder{
-				compiler: c,
-				edges:    edges,
-			}
-			Walk(vis, r)
-		}
-	}
+	c.RuleGraph = NewRuleGraph(c.Modules, c.GetRules)
 }
 
 type queryCompiler struct {
@@ -1106,40 +1093,145 @@ func (wc *withModifierChecker) err(code string, loc *Location, f string, a ...in
 	wc.errors = append(wc.errors, NewError(code, loc, f, a...))
 }
 
-type ruleGraphBuilder struct {
-	compiler *Compiler
-	edges    map[*Rule]struct{}
+// RuleGraph represents the dependencies between rules.
+type RuleGraph struct {
+	adj    map[*Rule]map[*Rule]struct{}
+	nodes  map[*Rule]struct{}
+	sorted []*Rule
 }
 
-func (vis *ruleGraphBuilder) Visit(v interface{}) Visitor {
-	ref, ok := v.(Ref)
+// NewRuleGraph returns a new RuleGraph based on modules. The list function
+// must return the rules referred to directly by the ref.
+func NewRuleGraph(modules map[string]*Module, list func(Ref) []*Rule) *RuleGraph {
+
+	ruleGraph := &RuleGraph{
+		adj:    map[*Rule]map[*Rule]struct{}{},
+		nodes:  map[*Rule]struct{}{},
+		sorted: nil,
+	}
+
+	// Walk over all rules, add them to graph, and build adjencency lists.
+	for _, module := range modules {
+		for _, ruleA := range module.Rules {
+			ruleGraph.addNode(ruleA)
+			WalkRefs(ruleA, func(ref Ref) bool {
+				for _, ruleB := range list(ref.GroundPrefix()) {
+					ruleGraph.addDependency(ruleA, ruleB)
+				}
+				return false
+			})
+		}
+	}
+
+	return ruleGraph
+}
+
+// Dependencies returns the set of rules that rule depends on.
+func (g *RuleGraph) Dependencies(rule *Rule) map[*Rule]struct{} {
+	return g.adj[rule]
+}
+
+// Sort returns a slice of rules sorted by dependencies. If a cycle is found,
+// ok is set to false.
+func (g *RuleGraph) Sort() (sorted []*Rule, ok bool) {
+
+	if g.sorted != nil {
+		return g.sorted, true
+	}
+
+	sort := &ruleGraphSort{
+		sorted: make([]*Rule, 0, len(g.nodes)),
+		deps:   g.Dependencies,
+		marked: map[*Rule]struct{}{},
+		temp:   map[*Rule]struct{}{},
+	}
+
+	for node := range g.nodes {
+		if !sort.Visit(node) {
+			return nil, false
+		}
+	}
+
+	g.sorted = sort.sorted
+	return g.sorted, true
+}
+
+func (g *RuleGraph) addDependency(u *Rule, v *Rule) {
+
+	if _, ok := g.nodes[u]; !ok {
+		g.addNode(u)
+	}
+
+	if _, ok := g.nodes[v]; !ok {
+		g.addNode(v)
+	}
+
+	edges, ok := g.adj[u]
 	if !ok {
-		return vis
+		edges = map[*Rule]struct{}{}
+		g.adj[u] = edges
 	}
 
-	for _, v := range vis.compiler.GetRules(ref.GroundPrefix()) {
-		vis.edges[v] = struct{}{}
-	}
-
-	return vis
+	edges[v] = struct{}{}
 }
 
-type ruleGraphTraveral struct {
-	graph   map[*Rule]map[*Rule]struct{}
+func (g *RuleGraph) addNode(n *Rule) {
+	g.nodes[n] = struct{}{}
+}
+
+type ruleGraphSort struct {
+	sorted []*Rule
+	deps   func(*Rule) map[*Rule]struct{}
+	marked map[*Rule]struct{}
+	temp   map[*Rule]struct{}
+}
+
+func (sort *ruleGraphSort) Marked(node *Rule) bool {
+	_, marked := sort.marked[node]
+	return marked
+}
+
+func (sort *ruleGraphSort) Visit(node *Rule) (ok bool) {
+	if _, ok := sort.temp[node]; ok {
+		return false
+	}
+	if sort.Marked(node) {
+		return true
+	}
+	sort.temp[node] = struct{}{}
+	for other := range sort.deps(node) {
+		if !sort.Visit(other) {
+			return false
+		}
+	}
+	sort.marked[node] = struct{}{}
+	delete(sort.temp, node)
+	sort.sorted = append(sort.sorted, node)
+	return true
+}
+
+type ruleGraphTraversal struct {
+	graph   *RuleGraph
 	visited map[*Rule]struct{}
 }
 
-func (g *ruleGraphTraveral) Edges(x util.T) []util.T {
+func newRuleGraphTraversal(graph *RuleGraph) *ruleGraphTraversal {
+	return &ruleGraphTraversal{
+		graph:   graph,
+		visited: map[*Rule]struct{}{},
+	}
+}
+
+func (g *ruleGraphTraversal) Edges(x util.T) []util.T {
 	u := x.(*Rule)
-	edges := g.graph[u]
 	r := []util.T{}
-	for v := range edges {
+	for v := range g.graph.Dependencies(u) {
 		r = append(r, v)
 	}
 	return r
 }
 
-func (g *ruleGraphTraveral) Visited(x util.T) bool {
+func (g *ruleGraphTraversal) Visited(x util.T) bool {
 	u := x.(*Rule)
 	_, ok := g.visited[u]
 	g.visited[u] = struct{}{}
diff --git a/ast/compile_test.go b/ast/compile_test.go
index 24fd8291df..07b79c47a6 100644
--- a/ast/compile_test.go
+++ b/ast/compile_test.go
@@ -488,15 +488,15 @@ p[foo[bar[i]]] = {"baz": baz} { true }`)
 	acTerm1 := ac(mod5.Rules[0])
 	assertTermEqual(t, acTerm1.Term, MustParseTerm("input.x.a"))
 	acTerm2 := ac(mod5.Rules[1])
-	assertTermEqual(t, acTerm2.Term, MustParseTerm("input.a.b.c.q.a"))
+	assertTermEqual(t, acTerm2.Term, MustParseTerm("data.a.b.c.q.a"))
 	acTerm3 := ac(mod5.Rules[2])
 	assertTermEqual(t, acTerm3.Body[0].Terms.([]*Term)[1], MustParseTerm("input.x.a"))
 	acTerm4 := ac(mod5.Rules[3])
-	assertTermEqual(t, acTerm4.Body[0].Terms.([]*Term)[1], MustParseTerm("input.a.b.c.q[i]"))
+	assertTermEqual(t, acTerm4.Body[0].Terms.([]*Term)[1], MustParseTerm("data.a.b.c.q[i]"))
 	acTerm5 := ac(mod5.Rules[4])
 	assertTermEqual(t, acTerm5.Body[0].Terms.([]*Term)[2].Value.(*ArrayComprehension).Term, MustParseTerm("input.x.a"))
 	acTerm6 := ac(mod5.Rules[5])
-	assertTermEqual(t, acTerm6.Body[0].Terms.([]*Term)[2].Value.(*ArrayComprehension).Body[0].Terms.([]*Term)[1], MustParseTerm("input.a.b.c.q[i]"))
+	assertTermEqual(t, acTerm6.Body[0].Terms.([]*Term)[2].Value.(*ArrayComprehension).Body[0].Terms.([]*Term)[1], MustParseTerm("data.a.b.c.q[i]"))
 
 	// Nested references.
 	mod6 := c.Modules["mod6"]
@@ -561,8 +561,83 @@ func TestCompilerSetRuleGraph(t *testing.T) {
 		r: struct{}{},
 	}
 
-	if !reflect.DeepEqual(edges, c.RuleGraph[p]) {
-		t.Errorf("Expected dependencies for p to be q and r but got: %v", c.RuleGraph[p])
+	if !reflect.DeepEqual(edges, c.RuleGraph.Dependencies(p)) {
+		t.Fatalf("Expected dependencies for p to be q and r but got: %v", c.RuleGraph.Dependencies(p))
+	}
+
+	sorted, ok := c.RuleGraph.Sort()
+	if !ok {
+		t.Fatalf("Expected sort to succeed.")
+	}
+
+	numRules := 0
+
+	for _, module := range c.Modules {
+		Walk(NewGenericVisitor(func(x interface{}) bool {
+			if _, ok := x.(*Rule); ok {
+				numRules++
+			}
+			return false
+		}), module)
+	}
+
+	if len(sorted) != numRules {
+		t.Fatalf("Expected numRules (%v) to be same as len(sorted) (%v)", numRules, len(sorted))
+	}
+
+	// Probe rules with dependencies. Ordering is not stable for ties because
+	// nodes are stored in a map.
+	probes := [][2]*Rule{
+		{c.Modules["mod1"].Rules[1], c.Modules["mod1"].Rules[0]}, // mod1.q before mod1.p
+		{c.Modules["mod2"].Rules[0], c.Modules["mod1"].Rules[0]}, // mod2.r before mod1.p
+		{c.Modules["mod1"].Rules[1], c.Modules["mod5"].Rules[1]}, // mod1.q before mod5.r
+		{c.Modules["mod1"].Rules[1], c.Modules["mod5"].Rules[3]}, // mod1.q before mod6.t
+		{c.Modules["mod1"].Rules[1], c.Modules["mod5"].Rules[5]}, // mod1.q before mod6.v
+		{c.Modules["mod6"].Rules[2], c.Modules["mod6"].Rules[3]}, // mod6.r before mod6.s
+	}
+
+	getSortedIdx := func(r *Rule) int {
+		for i := range sorted {
+			if sorted[i] == r {
+				return i
+			}
+		}
+		return -1
+	}
+
+	for num, probe := range probes {
+		i := getSortedIdx(probe[0])
+		j := getSortedIdx(probe[1])
+		if i == -1 || j == -1 {
+			t.Fatalf("Expected to find probe %d in sorted slice but got: i=%d, j=%d", num+1, i, j)
+		}
+		if i >= j {
+			t.Errorf("Sort order of probe %d (A) %v and (B) %v and is wrong (expected A before B)", num+1, probe[0], probe[1])
+		}
+	}
+}
+
+func TestRuleGraphCycle(t *testing.T) {
+	mod1 := `package a.b.c
+
+	p { q }
+	q { r }
+	r { s }
+	s { q }`
+
+	c := NewCompiler()
+	c.Modules = map[string]*Module{
+		"mod1": MustParseModule(mod1),
+	}
+
+	compileStages(c, "", "setRuleGraph")
+	assertNotFailed(t, c)
+
+	// fmt.Println(c.Modules)
+
+	_, ok := c.RuleGraph.Sort()
+	if ok {
+		t.Fatalf("Expected to find cycle in rule graph")
 	}
 
 }
@@ -1125,7 +1200,7 @@ x = false { true }`)
 	mod5 := MustParseModule(`package a.b.compr
 
 import input.x as y
-import input.a.b.c.q
+import data.a.b.c.q
 
 p = true { [y.a | true] }
 r = true { [q.a | true] }
diff --git a/util/graph_test.go b/util/graph_test.go
index 7d92618fe1..24fd140186 100644
--- a/util/graph_test.go
+++ b/util/graph_test.go
@@ -10,13 +10,24 @@ import (
 )
 
 type testTraversal struct {
-	g       map[int]map[int]struct{}
+	g       map[int][]int
 	visited map[int]struct{}
+	ordered []int
+	stop    *int
+}
+
+func newTestTraversal(g map[int][]int) *testTraversal {
+	return &testTraversal{
+		g:       g,
+		visited: map[int]struct{}{},
+		ordered: nil,
+		stop:    nil,
+	}
 }
 
 func (t *testTraversal) Edges(x T) []T {
 	r := []T{}
-	for v := range t.g[x.(int)] {
+	for _, v := range t.g[x.(int)] {
 		r = append(r, v)
 	}
 	return r
@@ -26,46 +37,137 @@ func (t *testTraversal) Equals(a, b T) bool {
 	return a.(int) == b.(int)
 }
 
+func (t *testTraversal) Iter(x T) bool {
+	t.ordered = append(t.ordered, x.(int))
+	return t.stop != nil && *t.stop == x.(int)
+}
+
 func (t *testTraversal) Visited(x T) bool {
 	_, ok := t.visited[x.(int)]
 	t.visited[x.(int)] = struct{}{}
 	return ok
 }
 
-func TestGraphDFS(t *testing.T) {
-
-	g := map[int]map[int]struct{}{
-		1: map[int]struct{}{
-			2: struct{}{},
-		},
-		2: map[int]struct{}{
-			3: struct{}{},
-			4: struct{}{},
-		},
-		3: map[int]struct{}{
-			2: struct{}{},
-		},
-		4: map[int]struct{}{
-			1: struct{}{},
-		},
+func TestDFSStop(t *testing.T) {
+	g := map[int][]int{
+		1: {2, 3},
+		2: {4, 5},
+		3: {6, 7},
+		6: {2},
 	}
 
-	t1 := &testTraversal{g, map[int]struct{}{}}
-	p1 := DFS(t1, 1, 2)
+	t1 := newTestTraversal(g)
+	stop := 6
+	t1.stop = &stop
+
+	stopped := DFS(t1, t1.Iter, 1)
+
+	if !stopped {
+		t.Fatalf("Expected DFS to stop but got: %v", t1.ordered)
+	}
+
+	expected := []int{1, 3, 7, 6}
+
+	if !reflect.DeepEqual(expected, t1.ordered) {
+		t.Fatalf("Expected DFS ordering %v but got: %v", expected, t1.ordered)
+	}
+}
+
+func TestBFSStop(t *testing.T) {
+	g := map[int][]int{
+		1: {2, 3},
+		2: {4, 5},
+		3: {6, 7},
+		6: {2},
+	}
+
+	t1 := newTestTraversal(g)
+	stop := 4
+	t1.stop = &stop
+
+	stopped := BFS(t1, t1.Iter, 1)
+
+	if !stopped {
+		t.Fatalf("Expected DFS to stop but got: %v", t1.ordered)
+	}
+
+	expected := []int{1, 2, 3, 4}
+
+	if !reflect.DeepEqual(expected, t1.ordered) {
+		t.Fatalf("Expected DFS ordering %v but got: %v", expected, t1.ordered)
+	}
+}
+
+func TestDFS(t *testing.T) {
+	g := map[int][]int{
+		1: {2, 3},
+		2: {4, 5},
+		3: {6, 7},
+		6: {2},
+	}
+
+	t1 := newTestTraversal(g)
+
+	stopped := DFS(t1, t1.Iter, 1)
+	if stopped {
+		t.Fatalf("Did not expect traversal to stop")
+	}
+
+	expected := []int{1, 3, 7, 6, 2, 5, 4}
+
+	if !reflect.DeepEqual(expected, t1.ordered) {
+		t.Fatalf("Expected DFS ordering %v but got: %v", expected, t1.ordered)
+	}
+}
+
+func TestBFS(t *testing.T) {
+	g := map[int][]int{
+		1: {2, 3},
+		2: {4, 5},
+		3: {6, 7},
+		6: {2},
+	}
+
+	t1 := newTestTraversal(g)
+
+	stopped := BFS(t1, t1.Iter, 1)
+	if stopped {
+		t.Fatalf("Did not expect traversal to stop")
+	}
+
+	expected := []int{1, 2, 3, 4, 5, 6, 7}
+
+	if !reflect.DeepEqual(expected, t1.ordered) {
+		t.Fatalf("Expected DFS ordering %v but got: %v", expected, t1.ordered)
+	}
+
+}
+
+func TestDFSPath(t *testing.T) {
+
+	g := map[int][]int{
+		1: {2},
+		2: {3, 4},
+		3: {2},
+		4: {1},
+	}
+
+	t1 := newTestTraversal(g)
+	p1 := DFSPath(t1, t1.Equals, 1, 2)
 
 	if !reflect.DeepEqual(p1, []T{1, 2}) {
 		t.Errorf("Expected DFS(1,2) to equal {1,2} but got: %v", p1)
 	}
 
-	t2 := &testTraversal{g, map[int]struct{}{}}
-	p2 := DFS(t2, 1, 4)
+	t2 := newTestTraversal(g)
+	p2 := DFSPath(t2, t2.Equals, 1, 4)
 
 	if !reflect.DeepEqual(p2, []T{1, 2, 4}) {
 		t.Errorf("Expected DFS(1,4) to equal {1,2,4} but got: %v", p2)
 	}
 
-	t3 := &testTraversal{g, map[int]struct{}{}}
-	p3 := DFS(t3, 1, 0xdeadbeef)
+	t3 := newTestTraversal(g)
+	p3 := DFSPath(t3, t3.Equals, 1, 0xdeadbeef)
 	if len(p3) != 0 {
 		t.Errorf("Expected DFS(1,0xdeadbeef to be empty but got: %v", p3)
 	}