topdown: use function arity to determine captured output (open-policy…

…-agent#3683)

Before, we had been passing `len(e.terms)` to `saveCall`, thereby
sidestepping the checks that would get the output variable into the
save set.

In the problematic policy,

    y = f(1)
    count(y)

with

    f(x) = [] { _ = input }

the call to `count(y)` would be mishandled: it would be evaluated,
since its arguments aren't in the saveset. Eval would fail because
its argument was a variable (see also
open-policy-agent#3680)

Fixes open-policy-agent#3681.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>
Signed-off-by: Dolev Farhi <farhi.dolev@gmail.com>

topdown/eval.go

@@ -1592,7 +1592,7 @@ func (e evalFunc) eval(iter unifyIterator) error {
 			unknown = e.e.unknown(e.terms[i], e.e.bindings)
 		}
 		if unknown {
-			return e.partialEvalSupport(iter)
+			return e.partialEvalSupport(argCount, iter)
 		}
 	}
 
@@ -1728,7 +1728,7 @@ func (e evalFunc) evalOneRule(iter unifyIterator, rule *ast.Rule, cacheKey ast.R
 	return result, err
 }
 
-func (e evalFunc) partialEvalSupport(iter unifyIterator) error {
+func (e evalFunc) partialEvalSupport(declArgsLen int, iter unifyIterator) error {
 
 	path := e.e.namespaceRef(e.ref)
 	term := ast.NewTerm(path)
@@ -1746,7 +1746,7 @@ func (e evalFunc) partialEvalSupport(iter unifyIterator) error {
 		return nil
 	}
 
-	return e.e.saveCall(len(e.terms), append([]*ast.Term{term}, e.terms[1:]...), iter)
+	return e.e.saveCall(declArgsLen, append([]*ast.Term{term}, e.terms[1:]...), iter)
 }
 
 func (e evalFunc) partialEvalSupportRule(rule *ast.Rule, path ast.Ref) error {

topdown/topdown_partial_test.go

@@ -2437,6 +2437,35 @@ func TestTopDownPartialEval(t *testing.T) {
 				`,
 			},
 		},
+		{
+			note:    "shallow inlining: functions with unknowns in body, result passed to builtin",
+			query:   "data.test.p",
+			shallow: true,
+			modules: []string{
+				`package test
+				p {
+				  y = f(1)
+				  count(y)
+				}
+
+				f(x) = [] {
+					# NOTE(sr): if we use '_' here, we cannot ever have a match
+					# when comparing the actual and expected support modules.
+					_x = input  # anything dependent on an unknown will do
+				}`,
+			},
+			wantQueries: []string{`data.partial.test.p = x_term_0_0; x_term_0_0`},
+			wantSupport: []string{
+				`package partial.test
+				p {
+					data.partial.test.f(1, __local1__1)
+					y1 = __local1__1
+					count(y1)
+				}
+				f(__local0__2) = [] { _x2 = input }
+				`,
+			},
+		},
 		{
 			note:  "comprehensions: ref heads (with namespacing)",
 			query: "data.test.p = true; input.x = x",