Fix evaluation bug with nested value refs

Nested references to array and object AST values were not being handled
correctly. Specifically, if the nested reference referred to a base
document, it was not being resolved. As a result, evaluation would stop
when the nested reference was encountered.

topdown/topdown.go

@@ -704,7 +704,7 @@ func (r resolver) Resolve(ref ast.Ref) (interface{}, error) {
 	return r.store.Read(r.context, r.txn, path)
 }
 
-// ResolveRefs returns the AST value obtained by resolving references to base
+// ResolveRefs returns the value obtained by resolving references to base
 // documents.
 func ResolveRefs(v ast.Value, t *Topdown) (ast.Value, error) {
 	result, err := ast.TransformRefs(v, func(r ast.Ref) (ast.Value, error) {
@@ -716,6 +716,18 @@ func ResolveRefs(v ast.Value, t *Topdown) (ast.Value, error) {
 	return result.(ast.Value), nil
 }
 
+// ResolveRefsTerm returns a copy of term obtained by resolving references to
+// base documents.
+func ResolveRefsTerm(term *ast.Term, t *Topdown) (*ast.Term, error) {
+	cpy := *term
+	var err error
+	cpy.Value, err = ResolveRefs(term.Value, t)
+	if err != nil {
+		return nil, err
+	}
+	return &cpy, nil
+}
+
 // ValueToInterface returns the underlying Go value associated with an AST value.
 // If the value is a reference, the reference is fetched from storage. Composite
 // AST values such as objects and arrays are converted recursively.
@@ -1686,6 +1698,14 @@ func evalRefRuleResultRec(t *Topdown, v ast.Value, ref, path ast.Ref, iter func(
 
 func evalRefRuleResultRecArray(t *Topdown, arr ast.Array, ref, path ast.Ref, iter func(*Topdown, ast.Value) error) error {
 	head, tail := ref[0], ref[1:]
+
+	if _, ok := head.Value.(ast.Ref); ok {
+		var err error
+		if head, err = ResolveRefsTerm(head, t); err != nil {
+			return err
+		}
+	}
+
 	switch n := head.Value.(type) {
 	case ast.Number:
 		idx, ok := n.Int()
@@ -1715,6 +1735,14 @@ func evalRefRuleResultRecArray(t *Topdown, arr ast.Array, ref, path ast.Ref, ite
 
 func evalRefRuleResultRecObject(t *Topdown, obj ast.Object, ref, path ast.Ref, iter func(*Topdown, ast.Value) error) error {
 	head, tail := ref[0], ref[1:]
+
+	if _, ok := head.Value.(ast.Ref); ok {
+		var err error
+		if head, err = ResolveRefsTerm(head, t); err != nil {
+			return err
+		}
+	}
+
 	switch k := head.Value.(type) {
 	case ast.String:
 		match := -1

topdown/topdown_test.go

@@ -835,8 +835,11 @@ func TestTopDownVarReferences(t *testing.T) {
 		{"embedded", []string{`p[x] { v = [1, 2, 3]; x = [{"a": v[i]}] }`}, `[[{"a": 1}], [{"a": 2}], [{"a": 3}]]`},
 		{"embedded ref binding", []string{`p[x] { v = c[i][j]; w = [v[0], v[1]]; x = w[y] }`}, "[null, false, true, 3.14159]"},
 		{"array: ground var", []string{`p[x] { i = [1, 2, 3, 4]; j = [1, 2, 999]; j[k] = y; i[y] = x }`}, "[2,3]"},
+		{"array: ref", []string{`p[y] { i = [1,2,3,4]; x = data.a[_]; i[x] = y }`}, `[2, 3, 4]`},
 		{"object: ground var", []string{`p[x] { i = {"a": 1, "b": 2, "c": 3}; j = ["a", "c", "deadbeef"]; j[k] = y; i[y] = x }`}, "[1, 3]"},
+		{"object: ref", []string{`p[y] { i = {"1": 1, "2": 2, "4": 4}; x = data.numbers[_]; i[x] = y }`}, `[1, 2, 4]`},
 		{"set: ground var", []string{`p[x] { i = {1, 2, 3, 4}; j = {1, 2, 99}; j[x]; i[x] }`}, "[1,2]"},
+		{"set: ref", []string{`p[x] { i = {1, 2, 3, 4}; x = data.a[_]; i[x] }`}, `[1, 2, 3, 4]`},
 		{"set: lookup: base docs", []string{`p = true { v = {[1, 999], [3, 4]}; pair = [a[2], 4]; v[pair] }`}, "true"},
 		{"set: lookup: embedded", []string{`p = true { x = [{}, {[1, 2], [3, 4]}]; y = [3, 4]; x[i][y] }`}, "true"},
 		{"set: lookup: dereference", []string{`p[[i, z, r]] { x = [{}, {[1, 2], [3, 4]}]; y = [3, 4]; x[i][y][z] = r }`}, "[[1,0,3], [1,1,4]]"},