ast: Fix compiler to expand exprs in rule args

Previously the compiler was not rewriting expr terms in the rule
args. In particular, this meant that indirect refs were not being
rewritten. This would lead to panics in the safety check which assume
that indirect refs have been rewritten (in other words, that reference
heads are always variables.)

This commit just updates the expr term rewriting to process rule args
as it does for the rule head key and value terms.

Fixes #2649

Signed-off-by: Torin Sandall <torinsandall@gmail.com>

ast/compile.go

@@ -3033,6 +3033,13 @@ func rewriteDynamicsComprehensionBody(original *Expr, f *equalityFactory, body B
 }
 
 func rewriteExprTermsInHead(gen *localVarGenerator, rule *Rule) {
+	for i := range rule.Head.Args {
+		support, output := expandExprTerm(gen, rule.Head.Args[i])
+		for j := range support {
+			rule.Body.Append(support[j])
+		}
+		rule.Head.Args[i] = output
+	}
 	if rule.Head.Key != nil {
 		support, output := expandExprTerm(gen, rule.Head.Key)
 		for i := range support {

ast/compile_test.go

@@ -1197,6 +1197,17 @@ func TestCompilerRewriteExprTerms(t *testing.T) {
 			`,
 			expected: Errors{&Error{Message: "rule arguments cannot contain calls"}},
 		},
+		{
+			note: "indirect ref in args",
+			module: `
+				package test
+
+				f([1][0]) { true }`,
+			expected: `
+				package test
+
+				f(__local0__[0]) { true; __local0__ = [1] }`,
+		},
 	}
 
 	for _, tc := range cases {