Allow indirect references

This allows the reference operator to work in more places, namely:

 -  After partial rules, e.g.:

        split("foo.bar", ".")[1]

 -  After array/set/object terms, e.g.:

        [1, 2, 3][i]
        {"foo": 1}.foo

 -  After array/set/object comprehensions, e.g.:

        concat(", ", ["Hello", name]) | {"Tom", "Jerry"}[name]][1]

This change was originally requested in open-policy-agent#1868.

Signed-off-by: Jasper Van der Jeugt <[email protected]>

ast/compile.go

@@ -2716,7 +2716,7 @@ func expandExprTerm(gen *localVarGenerator, term *Term) (support []*Expr, output
 		expr.Generated = true
 		support = append(support, expr)
 	case Ref:
-		support = expandExprTermSlice(gen, v)
+		support = expandExprRef(gen, v)
 	case Array:
 		support = expandExprTermSlice(gen, v)
 	case Object:
@@ -2765,6 +2765,35 @@ func expandExprTerm(gen *localVarGenerator, term *Term) (support []*Expr, output
 	return
 }
 
+func expandExprRef(gen *localVarGenerator, v []*Term) (support []*Expr) {
+	// Start by calling a normal expandExprTerm on all terms.
+	support = expandExprTermSlice(gen, v)
+
+	// Rewrite references in order to support indirect references.  We rewrite
+	// e.g.
+	//
+	//     [1, 2, 3][i]
+	//
+	// to
+	//
+	//     __local_var = [1, 2, 3]
+	//     __local_var[i]
+	//
+	// to support these.  This only impacts the reference subject, i.e. the
+	// first item in the slice.
+	if len(v) > 0 {
+		var subject = v[0]
+		switch subject.Value.(type) {
+		case Array, Object, Set, *ArrayComprehension, *SetComprehension, *ObjectComprehension, Call:
+			f := newEqualityFactory(gen)
+			assignToLocal := f.Generate(subject)
+			support = append(support, assignToLocal)
+			v[0] = assignToLocal.Operand(0)
+		}
+	}
+	return
+}
+
 func expandExprTermSlice(gen *localVarGenerator, v []*Term) (support []*Expr) {
 	for i := 0; i < len(v); i++ {
 		var extras []*Expr

ast/compile_test.go

@@ -893,6 +893,23 @@ func TestCompilerExprExpansion(t *testing.T) {
 				MustParseExpr("__local0__ = [[__local1__ | plus(z,1,__local2__); sum(y[__local2__], __local3__); eq(x, __local3__); plus(x, 1, __local1__)], __local4__]"),
 			},
 		},
+		{
+			note:  "indirect references",
+			input: `[1, 2, 3][i]`,
+			expected: []*Expr{
+				MustParseExpr("__local0__ = [1, 2, 3]"),
+				MustParseExpr("__local0__[i]"),
+			},
+		},
+		{
+			note:  "multiple indirect references",
+			input: `split(split("foo.bar:qux", ".")[_], ":")[i]`,
+			expected: []*Expr{
+				MustParseExpr(`split("foo.bar:qux", ".", __local0__)`),
+				MustParseExpr(`split(__local0__[_], ":", __local1__)`),
+				MustParseExpr(`__local1__[i]`),
+			},
+		},
 	}
 
 	for _, tc := range tests {