From d863e72d54e0507c24430e96fd97710c8b71f16b Mon Sep 17 00:00:00 2001
From: Torin Sandall <torinsandall@gmail.com>
Date: Fri, 13 Oct 2017 14:49:41 -0700
Subject: [PATCH] Fix wildcard mangling in rule head

Now that the AST walker can visit terms (as well as values) it's better
to just walk all terms in the statement and mangle them (rather than
relying on speical handling for collections of terms, which is bound to
break over time.)

Fixes #480
---
 ast/parser_ext.go  | 36 ++++++------------------------------
 ast/parser_test.go | 15 +++++++++++++--
 2 files changed, 19 insertions(+), 32 deletions(-)

diff --git a/ast/parser_ext.go b/ast/parser_ext.go
index 9ef856fd88..fb97c98ea7 100644
--- a/ast/parser_ext.go
+++ b/ast/parser_ext.go
@@ -609,40 +609,16 @@ type wildcardMangler struct {
 }
 
 func (vis *wildcardMangler) Visit(x interface{}) Visitor {
-	switch x := x.(type) {
-	case Object:
-		for _, i := range x {
-			vis.mangleSlice(i[:])
-		}
-	case Array:
-		vis.mangleSlice(x)
-	case *Set:
-		vis.mangleSlice(*x)
-	case Ref:
-		vis.mangleSlice(x)
-	case *Expr:
-		switch ts := x.Terms.(type) {
-		case []*Term:
-			vis.mangleSlice(ts)
-		case *Term:
-			vis.mangle(ts)
-		}
+	term, ok := x.(*Term)
+	if !ok {
+		return vis
 	}
-	return vis
-}
-
-func (vis *wildcardMangler) mangle(x *Term) {
-	if x.Equal(Wildcard) {
+	if term.Equal(Wildcard) {
 		name := fmt.Sprintf("%s%d", WildcardPrefix, vis.c)
-		x.Value = Var(name)
+		term.Value = Var(name)
 		vis.c++
 	}
-}
-
-func (vis *wildcardMangler) mangleSlice(xs []*Term) {
-	for _, x := range xs {
-		vis.mangle(x)
-	}
+	return vis
 }
 
 func setRuleModule(rule *Rule, module *Module) {
diff --git a/ast/parser_test.go b/ast/parser_test.go
index 726cb1e967..7a07ee398a 100644
--- a/ast/parser_test.go
+++ b/ast/parser_test.go
@@ -1194,9 +1194,9 @@ func TestWildcards(t *testing.T) {
 
 	assertParseOneTerm(t, "nested", `[{"a": a[_]}, _, {"b": _}]`, ArrayTerm(
 		ObjectTerm(
-			Item(StringTerm("a"), RefTerm(VarTerm("a"), VarTerm("$1"))),
+			Item(StringTerm("a"), RefTerm(VarTerm("a"), VarTerm("$0"))),
 		),
-		VarTerm("$0"),
+		VarTerm("$1"),
 		ObjectTerm(
 			Item(StringTerm("b"), VarTerm("$2")),
 		),
@@ -1219,6 +1219,17 @@ func TestWildcards(t *testing.T) {
 				),
 			),
 		)))
+
+	assertParseRule(t, "functions", `f(_) = y { true }`, &Rule{
+		Head: &Head{
+			Name: Var("f"),
+			Args: Args{
+				VarTerm("$0"),
+			},
+			Value: VarTerm("y"),
+		},
+		Body: NewBody(NewExpr(BooleanTerm(true))),
+	})
 }
 
 func TestRuleModulePtr(t *testing.T) {