From 831c34248a3608fda8c07768567489ddec0c8d6a Mon Sep 17 00:00:00 2001
From: Stephan Renatus <stephan.renatus@gmail.com>
Date: Mon, 26 Jul 2021 14:02:45 +0200
Subject: [PATCH] topdown/copypropagation: only update call bindings when
 output is captured

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 #3071.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>
---
 topdown/copypropagation/copypropagation.go | 10 ++++++----
 topdown/topdown_partial_test.go            | 15 +++++++++++++++
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/topdown/copypropagation/copypropagation.go b/topdown/copypropagation/copypropagation.go
index 2c325aeccc..2dee16f09c 100644
--- a/topdown/copypropagation/copypropagation.go
+++ b/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)
diff --git a/topdown/topdown_partial_test.go b/topdown/topdown_partial_test.go
index e3a27b1259..61aea119c7 100644
--- a/topdown/topdown_partial_test.go
+++ b/topdown/topdown_partial_test.go
@@ -2694,6 +2694,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()