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

@@ -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) {

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) {