topdown/copypropagation: only update call bindings when output is cap…

…tured

updateBindings before checked if the last operand was a variable, and attempted to
replace it if it was. With `neq(s & set(), s)`, the last operand was a variable,
but not the output variable -- the output wasn't captured.

So now, we use the arity of the called function to determine if we can replace a
call.

Fixes open-policy-agent#3071.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>

topdown/copypropagation/copypropagation.go

@@ -299,10 +299,12 @@ func (p *CopyPropagator) updateBindings(pctx *plugContext, expr *ast.Expr) bool
 		}
 	} else if expr.IsCall() {
 		terms := expr.Terms.([]*ast.Term)
-		output := terms[len(terms)-1]
-		if k, ok := output.Value.(ast.Var); ok && !p.livevars.Contains(k) && !pctx.headvars.Contains(k) {
-			pctx.removedEqs.Put(k, ast.CallTerm(terms[:len(terms)-1]...).Value)
-			return false
+		if p.compiler.GetArity(expr.Operator()) == len(terms)-2 { // with captured output
+			output := terms[len(terms)-1]
+			if k, ok := output.Value.(ast.Var); ok && !p.livevars.Contains(k) && !pctx.headvars.Contains(k) {
+				pctx.removedEqs.Put(k, ast.CallTerm(terms[:len(terms)-1]...).Value)
+				return false
+			}
 		}
 	}
 	return !isNoop(expr)

topdown/topdown_partial_test.go

@@ -2723,6 +2723,21 @@ func TestTopDownPartialEval(t *testing.T) {
 			},
 			wantQueries: []string{`x_term_1_01; x_term_1_01 = input[x_term_1_01]`},
 		},
+		{
+			note:  "copypropagation: circular reference (bug 3071)",
+			query: "data.test.p",
+			modules: []string{`package test
+				p[y] {
+					s := { i | input[i] }
+					s & set() != s
+					y := sprintf("%v", [s])
+				}`,
+			},
+			wantQueries: []string{`data.partial.test.p`},
+			wantSupport: []string{`package partial.test
+				p[__local1__1] { __local0__1 = {i1 | input[i1]}; neq(and(__local0__1, set()), __local0__1); sprintf("%v", [__local0__1], __local1__1) }
+			`},
+		},
 	}
 
 	ctx := context.Background()