gopls/internal/golang/extract: preserve comments in extracted block

Use printer.CommentedNode to preserve comments in function and method
extraction.

Fixes golang/go#50851

Change-Id: I7d8aa2683c980e613592f64646f8077952ea61be
Reviewed-on: https://go-review.googlesource.com/c/tools/+/629376
Reviewed-by: Alan Donovan <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: findleyr

gopls/internal/golang/extract.go

@@ -10,6 +10,7 @@ import (
 	"go/ast"
 	"go/format"
 	"go/parser"
+	"go/printer"
 	"go/token"
 	"go/types"
 	"slices"
@@ -449,7 +450,8 @@ func extractFunctionMethod(fset *token.FileSet, start, end token.Pos, src []byte
 		return nil, nil, err
 	}
 	selection := src[startOffset:endOffset]
-	extractedBlock, err := parseBlockStmt(fset, selection)
+
+	extractedBlock, extractedComments, err := parseStmts(fset, selection)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -570,26 +572,53 @@ func extractFunctionMethod(fset *token.FileSet, start, end token.Pos, src []byte
 	if canDefine {
 		sym = token.DEFINE
 	}
-	var name, funName string
+	var funName string
 	if isMethod {
-		name = "newMethod"
 		// TODO(suzmue): generate a name that does not conflict for "newMethod".
-		funName = name
+		funName = "newMethod"
 	} else {
-		name = "newFunction"
-		funName, _ = generateAvailableName(start, path, pkg, info, name, 0)
+		funName, _ = generateAvailableName(start, path, pkg, info, "newFunction", 0)
 	}
 	extractedFunCall := generateFuncCall(hasNonNestedReturn, hasReturnValues, params,
 		append(returns, getNames(retVars)...), funName, sym, receiverName)
 
-	// Build the extracted function.
+	// Create variable declarations for any identifiers that need to be initialized prior to
+	// calling the extracted function. We do not manually initialize variables if every return
+	// value is uninitialized. We can use := to initialize the variables in this situation.
+	var declarations []ast.Stmt
+	if canDefineCount != len(returns) {
+		declarations = initializeVars(uninitialized, retVars, seenUninitialized, seenVars)
+	}
+
+	var declBuf, replaceBuf, newFuncBuf, ifBuf, commentBuf bytes.Buffer
+	if err := format.Node(&declBuf, fset, declarations); err != nil {
+		return nil, nil, err
+	}
+	if err := format.Node(&replaceBuf, fset, extractedFunCall); err != nil {
+		return nil, nil, err
+	}
+	if ifReturn != nil {
+		if err := format.Node(&ifBuf, fset, ifReturn); err != nil {
+			return nil, nil, err
+		}
+	}
+
+	// Build the extracted function. We format the function declaration and body
+	// separately, so that comments are printed relative to the extracted
+	// BlockStmt.
+	//
+	// In other words, extractedBlock and extractedComments were parsed from a
+	// synthetic function declaration of the form func _() { ... }. If we now
+	// print the real function declaration, the length of the signature will have
+	// grown, causing some comment positions to be computed as inside the
+	// signature itself.
 	newFunc := &ast.FuncDecl{
 		Name: ast.NewIdent(funName),
 		Type: &ast.FuncType{
 			Params:  &ast.FieldList{List: paramTypes},
 			Results: &ast.FieldList{List: append(returnTypes, getDecls(retVars)...)},
 		},
-		Body: extractedBlock,
+		// Body handled separately -- see above.
 	}
 	if isMethod {
 		var names []*ast.Ident
@@ -603,39 +632,20 @@ func extractFunctionMethod(fset *token.FileSet, start, end token.Pos, src []byte
 			}},
 		}
 	}
-
-	// Create variable declarations for any identifiers that need to be initialized prior to
-	// calling the extracted function. We do not manually initialize variables if every return
-	// value is uninitialized. We can use := to initialize the variables in this situation.
-	var declarations []ast.Stmt
-	if canDefineCount != len(returns) {
-		declarations = initializeVars(uninitialized, retVars, seenUninitialized, seenVars)
-	}
-
-	var declBuf, replaceBuf, newFuncBuf, ifBuf, commentBuf bytes.Buffer
-	if err := format.Node(&declBuf, fset, declarations); err != nil {
+	if err := format.Node(&newFuncBuf, fset, newFunc); err != nil {
 		return nil, nil, err
 	}
-	if err := format.Node(&replaceBuf, fset, extractedFunCall); err != nil {
+	// Write a space between the end of the function signature and opening '{'.
+	if err := newFuncBuf.WriteByte(' '); err != nil {
 		return nil, nil, err
 	}
-	if ifReturn != nil {
-		if err := format.Node(&ifBuf, fset, ifReturn); err != nil {
-			return nil, nil, err
-		}
+	commentedNode := &printer.CommentedNode{
+		Node:     extractedBlock,
+		Comments: extractedComments,
 	}
-	if err := format.Node(&newFuncBuf, fset, newFunc); err != nil {
+	if err := format.Node(&newFuncBuf, fset, commentedNode); err != nil {
 		return nil, nil, err
 	}
-	// Find all the comments within the range and print them to be put somewhere.
-	// TODO(suzmue): print these in the extracted function at the correct place.
-	for _, cg := range file.Comments {
-		if cg.Pos().IsValid() && cg.Pos() < end && cg.Pos() >= start {
-			for _, c := range cg.List {
-				fmt.Fprintln(&commentBuf, c.Text)
-			}
-		}
-	}
 
 	// We're going to replace the whole enclosing function,
 	// so preserve the text before and after the selected block.
@@ -1187,25 +1197,25 @@ func varOverridden(info *types.Info, firstUse *ast.Ident, obj types.Object, isFr
 	return isOverriden
 }
 
-// parseBlockStmt generates an AST file from the given text. We then return the portion of the
-// file that represents the text.
-func parseBlockStmt(fset *token.FileSet, src []byte) (*ast.BlockStmt, error) {
+// parseStmts parses the specified source (a list of statements) and
+// returns them as a BlockStmt along with any associated comments.
+func parseStmts(fset *token.FileSet, src []byte) (*ast.BlockStmt, []*ast.CommentGroup, error) {
 	text := "package main\nfunc _() { " + string(src) + " }"
-	extract, err := parser.ParseFile(fset, "", text, parser.SkipObjectResolution)
+	file, err := parser.ParseFile(fset, "", text, parser.ParseComments|parser.SkipObjectResolution)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
-	if len(extract.Decls) == 0 {
-		return nil, fmt.Errorf("parsed file does not contain any declarations")
+	if len(file.Decls) != 1 {
+		return nil, nil, fmt.Errorf("got %d declarations, want 1", len(file.Decls))
 	}
-	decl, ok := extract.Decls[0].(*ast.FuncDecl)
+	decl, ok := file.Decls[0].(*ast.FuncDecl)
 	if !ok {
-		return nil, fmt.Errorf("parsed file does not contain expected function declaration")
+		return nil, nil, bug.Errorf("parsed file does not contain expected function declaration")
 	}
 	if decl.Body == nil {
-		return nil, fmt.Errorf("extracted function has no body")
+		return nil, nil, bug.Errorf("extracted function has no body")
 	}
-	return decl.Body, nil
+	return decl.Body, file.Comments, nil
 }
 
 // generateReturnInfo generates the information we need to adjust the return statements and

gopls/internal/test/marker/testdata/codeaction/extract_method.txt

@@ -237,20 +237,14 @@ func (b *B) LongListWithT(ctx context.Context, t *testing.T) (int, error) {
 +}
 +
 -- @contextFuncB/context.go --
-@@ -33 +33,6 @@
--	sum := b.x + b.y //@loc(B_AddPWithB, re`(?s:^.*?Err\(\))`)
-+	//@loc(B_AddPWithB, re`(?s:^.*?Err\(\))`)
+@@ -33 +33,4 @@
 +	return newFunction(ctx, tB, b)
 +}
 +
 +func newFunction(ctx context.Context, tB *testing.B, b *B) (int, error) {
-+	sum := b.x + b.y
 -- @contextFuncT/context.go --
-@@ -42 +42,6 @@
--	p4 := p1 + p2 //@loc(B_LongListWithT, re`(?s:^.*?Err\(\))`)
-+	//@loc(B_LongListWithT, re`(?s:^.*?Err\(\))`)
+@@ -42 +42,4 @@
 +	return newFunction(ctx, t, p1, p2, p3)
 +}
 +
 +func newFunction(ctx context.Context, t *testing.T, p1 int, p2 int, p3 int) (int, error) {
-+	p4 := p1 + p2