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 open-policy-agent#490

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)

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) {