go/ast/astutil: PathEnclosingInterval: add missing FuncType nodes

The special case for FuncDecl > FuncType added long ago had the
effect of eliding FuncType from the returned path. Add it back.

Fixes golang/go#68202

Change-Id: I7dd1bc5c81c89af9b1f388154bec4e910d0fb2d6
Reviewed-on: https://go-review.googlesource.com/c/tools/+/595195
Reviewed-by: Robert Findley <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

go/ast/astutil/enclosing.go

@@ -106,8 +106,21 @@ func PathEnclosingInterval(root *ast.File, start, end token.Pos) (path []ast.Nod
 
 			// Does augmented child strictly contain [start, end)?
 			if augPos <= start && end <= augEnd {
-				_, isToken := child.(tokenNode)
-				return isToken || visit(child)
+				if is[tokenNode](child) {
+					return true
+				}
+
+				// childrenOf elides the FuncType node beneath FuncDecl.
+				// Add it back here for TypeParams, Params, Results,
+				// all FieldLists). But we don't add it back for the "func" token
+				// even though it is is the tree at FuncDecl.Type.Func.
+				if decl, ok := node.(*ast.FuncDecl); ok {
+					if fields, ok := child.(*ast.FieldList); ok && fields != decl.Recv {
+						path = append(path, decl.Type)
+					}
+				}
+
+				return visit(child)
 			}
 
 			// Does [start, end) overlap multiple children?
@@ -313,6 +326,8 @@ func childrenOf(n ast.Node) []ast.Node {
 		//
 		// As a workaround, we inline the case for FuncType
 		// here and order things correctly.
+		// We also need to insert the elided FuncType just
+		// before the 'visit' recursion.
 		//
 		children = nil // discard ast.Walk(FuncDecl) info subtrees
 		children = append(children, tok(n.Type.Func, len("func")))
@@ -632,3 +647,8 @@ func NodeDescription(n ast.Node) string {
 	}
 	panic(fmt.Sprintf("unexpected node type: %T", n))
 }
+
+func is[T any](x any) bool {
+	_, ok := x.(T)
+	return ok
+}

go/ast/astutil/enclosing_test.go

@@ -72,6 +72,8 @@ func g[A any, P interface{ctype1| ~ctype2}](a1 A, p1 P) {}
 
 type PT[T constraint] struct{ t T }
 
+func (r recv) method(p param) {}
+
 var v GT[targ1]
 
 var h = g[ targ2, targ3]
@@ -204,14 +206,16 @@ func TestPathEnclosingInterval_Paths(t *testing.T) {
 			"[Ident File],true"},
 		{"f() // NB",
 			"[CallExpr ExprStmt BlockStmt FuncDecl File],true"},
-		{" any", "[Ident Field FieldList FuncDecl File],true"},
-		{"|", "[BinaryExpr Field FieldList InterfaceType Field FieldList FuncDecl File],true"},
+		{" any", "[Ident Field FieldList FuncType FuncDecl File],true"},
+		{"|", "[BinaryExpr Field FieldList InterfaceType Field FieldList FuncType FuncDecl File],true"},
 		{"ctype2",
-			"[Ident UnaryExpr BinaryExpr Field FieldList InterfaceType Field FieldList FuncDecl File],true"},
-		{"a1", "[Ident Field FieldList FuncDecl File],true"},
+			"[Ident UnaryExpr BinaryExpr Field FieldList InterfaceType Field FieldList FuncType FuncDecl File],true"},
+		{"a1", "[Ident Field FieldList FuncType FuncDecl File],true"},
 		{"PT[T constraint]", "[TypeSpec GenDecl File],false"},
 		{"[T constraint]", "[FieldList TypeSpec GenDecl File],true"},
 		{"targ2", "[Ident IndexListExpr ValueSpec GenDecl File],true"},
+		{"p param", "[Field FieldList FuncType FuncDecl File],true"}, // FuncType is present for FuncDecl.Params (etc)
+		{"r recv", "[Field FieldList FuncDecl File],true"},           // no FuncType for FuncDecl.Recv
 	}
 	for _, test := range tests {
 		f, start, end := findInterval(t, new(token.FileSet), input, test.substr)