diff --git a/rego/rego.go b/rego/rego.go index 6601b42fcd..d0a09f24e4 100644 --- a/rego/rego.go +++ b/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) diff --git a/rego/rego_test.go b/rego/rego_test.go index 4e5ce3b764..1f669122d1 100644 --- a/rego/rego_test.go +++ b/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{ diff --git a/types/types.go b/types/types.go index fdfc1f879c..256efcc05c 100644 --- a/types/types.go +++ b/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 {