From a09615c753003bd85494bfbfdd98c69fd7700624 Mon Sep 17 00:00:00 2001
From: Stephan Renatus <stephan.renatus@gmail.com>
Date: Tue, 29 Nov 2022 13:49:35 +0100
Subject: [PATCH] ast/parser: fresh wildcards for chained function heads

This is a spiritual follow-up to #5412.

A policy like

    f(_) := 1 { true } { true } { true }

would have been pretty-printed as

    f(_0) := 1
    f(_0) := 1
    f(_0) := 1

because of the duplicated wildcards.

They now get the same treatment as "else" gets: fresh wildcards.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>
---
 ast/parser.go      |  9 ++++++---
 ast/parser_test.go | 20 ++++++++++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/ast/parser.go b/ast/parser.go
index 0f9957f129..0b523d372c 100644
--- a/ast/parser.go
+++ b/ast/parser.go
@@ -660,9 +660,7 @@ func (p *Parser) parseRules() []*Rule {
 
 	rule.Location.Text = p.s.Text(rule.Location.Offset, p.s.lastEnd)
 
-	var rules []*Rule
-
-	rules = append(rules, &rule)
+	rules := []*Rule{&rule}
 
 	for p.s.tok == tokens.LBrace {
 
@@ -688,6 +686,11 @@ func (p *Parser) parseRules() []*Rule {
 		// rule's head AST but have their location
 		// set to the rule body.
 		next.Head = rule.Head.Copy()
+		for i := range next.Head.Args {
+			if v, ok := next.Head.Args[i].Value.(Var); ok && v.IsWildcard() {
+				next.Head.Args[i].Value = Var(p.genwildcard())
+			}
+		}
 		setLocRecursive(next.Head, loc)
 
 		rules = append(rules, &next)
diff --git a/ast/parser_test.go b/ast/parser_test.go
index d7500dc1b9..4568ac3b09 100644
--- a/ast/parser_test.go
+++ b/ast/parser_test.go
@@ -1658,6 +1658,26 @@ func TestRule(t *testing.T) {
 			Body: MustParseBody(`true`),
 		},
 	})
+
+	name := Var("f")
+	ref := Ref{VarTerm("f")}
+	tr := BooleanTerm(true)
+	head := func(v string) *Head { return &Head{Name: name, Reference: ref, Value: tr, Args: []*Term{VarTerm(v)}} }
+	assertParseModule(t, "wildcard in chained function heads", `package test
+	f(_) { true } { true }
+`, &Module{
+		Package: MustParsePackage(`package test`),
+		Rules: []*Rule{
+			{
+				Head: head("$0"),
+				Body: MustParseBody("true"),
+			},
+			{
+				Head: head("$1"),
+				Body: MustParseBody("true"),
+			},
+		},
+	})
 }
 
 func TestRuleContains(t *testing.T) {