From 892ee15608bfeb5aa5bedacd2be1f05ea84374fd Mon Sep 17 00:00:00 2001 From: Stephan Renatus Date: Thu, 2 Dec 2021 13:52:16 +0100 Subject: [PATCH] ast/compile: fix print rewriting for arrays in := and head vars For both partial and complete rules, when variables from destructured arrays in assignments had been used in the rule head, the print rewriting would fail. Now, this situation is accounted for by starting the safe VarSet with r.Head.Vars() instead of r.Head.Args.Vars() The former will check each of the Args, Key and Value fields of Head: if it's non-empty, it'll add .Vars() to the returned VarSet. Fixes #4078. Signed-off-by: Stephan Renatus --- ast/compile.go | 2 +- ast/compile_test.go | 30 ++++++++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/ast/compile.go b/ast/compile.go index f446606d16..983af7374e 100644 --- a/ast/compile.go +++ b/ast/compile.go @@ -1371,7 +1371,7 @@ func (c *Compiler) rewritePrintCalls() { for _, name := range c.sorted { mod := c.Modules[name] WalkRules(mod, func(r *Rule) bool { - safe := r.Head.Args.Vars() + safe := r.Head.Vars() safe.Update(ReservedVars) WalkBodies(r, func(b Body) bool { for _, err := range rewritePrintCalls(c.localvargen, c.GetArity, safe, b) { diff --git a/ast/compile_test.go b/ast/compile_test.go index dce2baf216..223f7eac7b 100644 --- a/ast/compile_test.go +++ b/ast/compile_test.go @@ -3169,6 +3169,36 @@ func TestCompilerRewritePrintCalls(t *testing.T) { f(__local0__) = __local2__ { true; __local2__ = {1 | __local0__[x]; __local3__ = {__local1__ | __local1__ = x}; internal.print([__local3__])} } `, }, + { + note: "print call of var in head key", + module: `package test + f(_) = [1, 2, 3] + p[x] { [_, x, _] := f(true); print(x) }`, + exp: `package test + f(__local0__) = [1, 2, 3] { true } + p[__local2__] { data.test.f(true, __local5__); [__local1__, __local2__, __local3__] = __local5__; __local6__ = {__local4__ | __local4__ = __local2__}; internal.print([__local6__]) } + `, + }, + { + note: "print call of var in head value", + module: `package test + f(_) = [1, 2, 3] + p = x { [_, x, _] := f(true); print(x) }`, + exp: `package test + f(__local0__) = [1, 2, 3] { true } + p = __local2__ { data.test.f(true, __local5__); [__local1__, __local2__, __local3__] = __local5__; __local6__ = {__local4__ | __local4__ = __local2__}; internal.print([__local6__]) } + `, + }, + { + note: "print call of vars in head key and value", + module: `package test + f(_) = [1, 2, 3] + p[x] = y { [_, x, y] := f(true); print(x) }`, + exp: `package test + f(__local0__) = [1, 2, 3] { true } + p[__local2__] = __local3__ { data.test.f(true, __local5__); [__local1__, __local2__, __local3__] = __local5__; __local6__ = {__local4__ | __local4__ = __local2__}; internal.print([__local6__]) } + `, + }, } for _, tc := range cases {