Fix with value rewriting for call terms

If a call term was used as a with value, e.g., "p with input as f(1)"
the type checker would panic because the call was not expanded as
expected. These changes update the rewrite term expr stage to rewrite
with values like other AST nodes such as calls and composites.

Also, add an extra test case for with value rewriting to cover open-policy-agent#772.

Fixes open-policy-agent#916

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

ast/compile.go

@@ -2277,6 +2277,11 @@ func rewriteExprTermsInBody(gen *localVarGenerator, body Body) Body {
 }
 
 func expandExpr(gen *localVarGenerator, expr *Expr) (result []*Expr) {
+	for i := range expr.With {
+		extras, value := expandExprTerm(gen, expr.With[i].Value)
+		expr.With[i].Value = value
+		result = append(result, extras...)
+	}
 	switch terms := expr.Terms.(type) {
 	case *Term:
 		extras, term := expandExprTerm(gen, terms)

ast/compile_test.go

@@ -812,6 +812,8 @@ func TestCompilerRewriteExprTerms(t *testing.T) {
 		f(x) = data.test.g(x)
 
 		pi = 3 + .14
+
+		with_value { 1 with input as f(1) }
 	`
 
 	compiler := NewCompiler()
@@ -833,6 +835,8 @@ func TestCompilerRewriteExprTerms(t *testing.T) {
 		f(x) = __local4__ { true; data.test.g(x, __local4__) }
 
 		pi = __local5__ { true; plus(3, 0.14, __local5__) }
+
+		with_value { data.test.f(1, __local6__); 1 with input as __local6__ }
 	`)
 
 	if !expected.Equal(compiler.Modules["test"]) {
@@ -2121,6 +2125,7 @@ func TestQueryCompiler(t *testing.T) {
 		{"safe vars", `data; abc`, `package ex`, []string{"import input.xyz as abc"}, `{}`, `data; input.xyz`},
 		{"reorder", `x != 1; x = 0`, "", nil, "", `x = 0; x != 1`},
 		{"bad with target", "x = 1 with data.p as null", "", nil, "", fmt.Errorf("1 error occurred: 1:12: rego_type_error: with keyword target must be input")},
+		{"rewrite with value", `1 with input as [z]`, "package a.b.c", nil, "", `__local1__ = data.a.b.c.z; __local0__ = [__local1__]; 1 with input as __local0__`},
 		{"unsafe exprs", "count(sum())", "", nil, "", fmt.Errorf("1 error occurred: 1:1: rego_unsafe_var_error: expression is unsafe")},
 		{"check types", "x = data.a.b.c.z; y = null; x = y", "", nil, "", fmt.Errorf("match error\n\tleft  : number\n\tright : null")},
 	}