go/ast/astutil: make Apply follow TypeParams fields

It completely ignored the TypeParams fields in FuncType and TypeSpec.
I was confused why one of my tools, which looks at *ast.Ident nodes,
saw the identifiers where type parameters were being used,
but not the identifiers where they were being declared.

We use the ForFuncType and ForTypeSpec helpers, just like in the rest of
the package, for consistency and backwards compatibility with Go 1.17.

Add a regression test as well.

Change-Id: I00b2dfe5e04d92d63e6d5e91c6466598c865ed0b
Reviewed-on: https://go-review.googlesource.com/c/tools/+/405194
Run-TryBot: Daniel Martí <mvdan@mvdan.cc>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Robert Findley <rfindley@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
diff --git a/go/ast/astutil/rewrite.go b/go/ast/astutil/rewrite.go
index 729e9c8..f430b21 100644
--- a/go/ast/astutil/rewrite.go
+++ b/go/ast/astutil/rewrite.go
@@ -293,6 +293,9 @@
 		a.apply(n, "Fields", nil, n.Fields)
 
 	case *ast.FuncType:
+		if tparams := typeparams.ForFuncType(n); tparams != nil {
+			a.apply(n, "TypeParams", nil, tparams)
+		}
 		a.apply(n, "Params", nil, n.Params)
 		a.apply(n, "Results", nil, n.Results)
 
@@ -405,6 +408,9 @@
 	case *ast.TypeSpec:
 		a.apply(n, "Doc", nil, n.Doc)
 		a.apply(n, "Name", nil, n.Name)
+		if tparams := typeparams.ForTypeSpec(n); tparams != nil {
+			a.apply(n, "TypeParams", nil, tparams)
+		}
 		a.apply(n, "Type", nil, n.Type)
 		a.apply(n, "Comment", nil, n.Comment)
 
diff --git a/go/ast/astutil/rewrite_test.go b/go/ast/astutil/rewrite_test.go
index 9d23170..4ef6fe9 100644
--- a/go/ast/astutil/rewrite_test.go
+++ b/go/ast/astutil/rewrite_test.go
@@ -202,20 +202,30 @@
 type T[P1, P2 any] int
 
 type R T[int, string]
+
+func F[Q1 any](q Q1) {}
 `,
+			// TODO: note how the rewrite adds a trailing comma in "func F".
+			// Is that a bug in the test, or in astutil.Apply?
 			want: `package p
 
-type S[P1, P2 any] int32
+type S[R1, P2 any] int32
 
 type R S[int32, string]
+
+func F[X1 any](q X1,) {}
 `,
 			post: func(c *astutil.Cursor) bool {
 				if ident, ok := c.Node().(*ast.Ident); ok {
-					if ident.Name == "int" {
+					switch ident.Name {
+					case "int":
 						c.Replace(ast.NewIdent("int32"))
-					}
-					if ident.Name == "T" {
+					case "T":
 						c.Replace(ast.NewIdent("S"))
+					case "P1":
+						c.Replace(ast.NewIdent("R1"))
+					case "Q1":
+						c.Replace(ast.NewIdent("X1"))
 					}
 				}
 				return true