rego: Do not rewrite/capture void function call return values

This commit prevents the rewriting step from attempting to capture the
result of void function calls. This avoids generating invalid queries
that will fail to type check.

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

rego/rego.go

@@ -2167,7 +2167,8 @@ func (r *Rego) rewriteQueryToCaptureValue(qc ast.QueryCompiler, query ast.Body)
 			expr.Terms = ast.Equality.Expr(terms, capture).Terms
 			r.capture[expr] = capture.Value.(ast.Var)
 		case []*ast.Term:
-			if r.compiler.GetArity(expr.Operator()) == len(terms)-1 {
+			tpe := r.compiler.TypeEnv.Get(terms[0])
+			if !types.Void(tpe) && types.Arity(tpe) == len(terms)-1 {
 				capture = r.generateTermVar()
 				expr.Terms = append(terms, capture)
 				r.capture[expr] = capture.Value.(ast.Var)

rego/rego_test.go

@@ -216,6 +216,22 @@ func TestRegoRewrittenVarsCapture(t *testing.T) {
 
 }
 
+func TestRegoDoNotCaptureVoidCalls(t *testing.T) {
+
+	ctx := context.Background()
+
+	r := New(Query("print(1)"))
+
+	rs, err := r.Eval(ctx)
+	if err != nil || len(rs) != 1 {
+		t.Fatal(err, "rs:", rs)
+	}
+
+	if !rs[0].Expressions[0].Value.(bool) {
+		t.Fatal("expected expression value to be true")
+	}
+}
+
 func TestRegoCancellation(t *testing.T) {
 
 	ast.RegisterBuiltin(&ast.Builtin{

types/types.go

@@ -467,6 +467,23 @@ func Args(x ...Type) []Type {
 	return x
 }
 
+// Void returns true if the function has no return value. This function returns
+// false if tpe is not a function.
+func Void(x Type) bool {
+	f, ok := x.(*Function)
+	return ok && f.Result() == nil
+}
+
+// Arity returns the number of arguments in the function signature. This
+// function returns -1 if tpe is not a function.
+func Arity(x Type) int {
+	f, ok := x.(*Function)
+	if !ok {
+		return -1
+	}
+	return len(f.FuncArgs().Args)
+}
+
 // NewFunction returns a new Function object where xs[:len(xs)-1] are arguments
 // and xs[len(xs)-1] is the result type.
 func NewFunction(args []Type, result Type) *Function {