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>

go/ast/astutil/rewrite.go

@@ -293,6 +293,9 @@ func (a *application) apply(parent ast.Node, name string, iter *iterator, n ast.
 		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 @@ func (a *application) apply(parent ast.Node, name string, iter *iterator, n ast.
 	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)
 

go/ast/astutil/rewrite_test.go

@@ -202,20 +202,30 @@ func init() {
 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