From bb4e94ab7853a2428b947561faa35bd161ddfaf0 Mon Sep 17 00:00:00 2001
From: Torin Sandall <torinsandall@gmail.com>
Date: Mon, 30 Oct 2017 10:07:41 -0700
Subject: [PATCH] Fix index usage for virtual docs

Topdown was trying to use the data index when iterating over modules.
This change modifies how topdown uses the rule tree to determine if the
reference refers to one or more rules.

Fixes #490
---
 topdown/topdown.go      | 23 ++++++++++-------------
 topdown/topdown_test.go | 12 +++++++++++-
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/topdown/topdown.go b/topdown/topdown.go
index 0a32adb1ba..ca80b5b709 100644
--- a/topdown/topdown.go
+++ b/topdown/topdown.go
@@ -2214,23 +2214,20 @@ func indexBuildLazy(t *Topdown, ref ast.Ref) (storage.Index, error) {
 	}
 
 	// Ignore refs against virtual docs.
-	tmp := ast.Ref{ref[0], ref[1]}
-	r := t.Compiler.GetRulesExact(tmp)
-	if r != nil {
-		return nil, nil
-	}
-
-	for _, p := range ref[2:] {
+	node := t.Compiler.RuleTree.Child(ref[0].Value)
 
-		if !p.Value.IsGround() {
+	for i := 1; i < len(ref); i++ {
+		if node == nil || !ref[i].IsGround() {
 			break
 		}
-
-		tmp = append(tmp, p)
-		r := t.Compiler.GetRulesExact(tmp)
-		if r != nil {
-			return nil, nil
+		if len(node.Values) > 0 {
+			break
 		}
+		node = node.Child(ref[i].Value)
+	}
+
+	if node != nil {
+		return nil, nil
 	}
 
 	index, err := t.Store.Build(t.Context, t.txn, ref)
diff --git a/topdown/topdown_test.go b/topdown/topdown_test.go
index 26a21bb0d0..2730fd84c7 100644
--- a/topdown/topdown_test.go
+++ b/topdown/topdown_test.go
@@ -777,6 +777,12 @@ p = true { input.foo }`,
 
 p = true { false }`,
 
+		`package topdown.virtual.constants
+
+		p = 1
+		q = 2
+		r = 1`,
+
 		// Define virtual docs that we can query to obtain merged result.
 		`package topdown
 
@@ -787,7 +793,10 @@ s = data.topdown.no { true }
 t = data.topdown.a.b.c.undefined1 { true }
 u = data.topdown.missing.input.value { true }
 v = data.topdown.g { true }
-w = data.topdown.set { true }`,
+w = data.topdown.set { true }
+
+iterate_ground[x] { data.topdown.virtual.constants[x] = 1 }
+`,
 	})
 
 	store := inmem.NewFromObject(data)
@@ -835,6 +844,7 @@ w = data.topdown.set { true }`,
 	assertTopDownWithPath(t, compiler, store, "base/virtual: undefined", []string{"topdown", "t"}, "{}", "{}")
 	assertTopDownWithPath(t, compiler, store, "base/virtual: undefined-2", []string{"topdown", "v"}, "{}", `{"h": {"k": [1,2,3]}}`)
 	assertTopDownWithPath(t, compiler, store, "base/virtual: missing input value", []string{"topdown", "u"}, "{}", "")
+	assertTopDownWithPath(t, compiler, store, "iterate ground", []string{"topdown", "iterate_ground"}, "{}", `["p", "r"]`)
 }
 
 func TestTopDownNestedReferences(t *testing.T) {