Skip to content

Commit 9dff42e

Browse files
committed
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]>
1 parent 8c3ba8c commit 9dff42e

File tree

6 files changed

+113
-96
lines changed

6 files changed

+113
-96
lines changed

gopls/internal/golang/extract.go

Lines changed: 55 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"go/ast"
1111
"go/format"
1212
"go/parser"
13+
"go/printer"
1314
"go/token"
1415
"go/types"
1516
"slices"
@@ -449,7 +450,8 @@ func extractFunctionMethod(fset *token.FileSet, start, end token.Pos, src []byte
449450
return nil, nil, err
450451
}
451452
selection := src[startOffset:endOffset]
452-
extractedBlock, err := parseBlockStmt(fset, selection)
453+
454+
extractedBlock, extractedComments, err := parseStmts(fset, selection)
453455
if err != nil {
454456
return nil, nil, err
455457
}
@@ -570,26 +572,53 @@ func extractFunctionMethod(fset *token.FileSet, start, end token.Pos, src []byte
570572
if canDefine {
571573
sym = token.DEFINE
572574
}
573-
var name, funName string
575+
var funName string
574576
if isMethod {
575-
name = "newMethod"
576577
// TODO(suzmue): generate a name that does not conflict for "newMethod".
577-
funName = name
578+
funName = "newMethod"
578579
} else {
579-
name = "newFunction"
580-
funName, _ = generateAvailableName(start, path, pkg, info, name, 0)
580+
funName, _ = generateAvailableName(start, path, pkg, info, "newFunction", 0)
581581
}
582582
extractedFunCall := generateFuncCall(hasNonNestedReturn, hasReturnValues, params,
583583
append(returns, getNames(retVars)...), funName, sym, receiverName)
584584

585-
// Build the extracted function.
585+
// Create variable declarations for any identifiers that need to be initialized prior to
586+
// calling the extracted function. We do not manually initialize variables if every return
587+
// value is uninitialized. We can use := to initialize the variables in this situation.
588+
var declarations []ast.Stmt
589+
if canDefineCount != len(returns) {
590+
declarations = initializeVars(uninitialized, retVars, seenUninitialized, seenVars)
591+
}
592+
593+
var declBuf, replaceBuf, newFuncBuf, ifBuf, commentBuf bytes.Buffer
594+
if err := format.Node(&declBuf, fset, declarations); err != nil {
595+
return nil, nil, err
596+
}
597+
if err := format.Node(&replaceBuf, fset, extractedFunCall); err != nil {
598+
return nil, nil, err
599+
}
600+
if ifReturn != nil {
601+
if err := format.Node(&ifBuf, fset, ifReturn); err != nil {
602+
return nil, nil, err
603+
}
604+
}
605+
606+
// Build the extracted function. We format the function declaration and body
607+
// separately, so that comments are printed relative to the extracted
608+
// BlockStmt.
609+
//
610+
// In other words, extractedBlock and extractedComments were parsed from a
611+
// synthetic function declaration of the form func _() { ... }. If we now
612+
// print the real function declaration, the length of the signature will have
613+
// grown, causing some comment positions to be computed as inside the
614+
// signature itself.
586615
newFunc := &ast.FuncDecl{
587616
Name: ast.NewIdent(funName),
588617
Type: &ast.FuncType{
589618
Params: &ast.FieldList{List: paramTypes},
590619
Results: &ast.FieldList{List: append(returnTypes, getDecls(retVars)...)},
591620
},
592-
Body: extractedBlock,
621+
// Body handled separately -- see above.
593622
}
594623
if isMethod {
595624
var names []*ast.Ident
@@ -603,39 +632,20 @@ func extractFunctionMethod(fset *token.FileSet, start, end token.Pos, src []byte
603632
}},
604633
}
605634
}
606-
607-
// Create variable declarations for any identifiers that need to be initialized prior to
608-
// calling the extracted function. We do not manually initialize variables if every return
609-
// value is uninitialized. We can use := to initialize the variables in this situation.
610-
var declarations []ast.Stmt
611-
if canDefineCount != len(returns) {
612-
declarations = initializeVars(uninitialized, retVars, seenUninitialized, seenVars)
613-
}
614-
615-
var declBuf, replaceBuf, newFuncBuf, ifBuf, commentBuf bytes.Buffer
616-
if err := format.Node(&declBuf, fset, declarations); err != nil {
635+
if err := format.Node(&newFuncBuf, fset, newFunc); err != nil {
617636
return nil, nil, err
618637
}
619-
if err := format.Node(&replaceBuf, fset, extractedFunCall); err != nil {
638+
// Write a space between the end of the function signature and opening '{'.
639+
if err := newFuncBuf.WriteByte(' '); err != nil {
620640
return nil, nil, err
621641
}
622-
if ifReturn != nil {
623-
if err := format.Node(&ifBuf, fset, ifReturn); err != nil {
624-
return nil, nil, err
625-
}
642+
commentedNode := &printer.CommentedNode{
643+
Node: extractedBlock,
644+
Comments: extractedComments,
626645
}
627-
if err := format.Node(&newFuncBuf, fset, newFunc); err != nil {
646+
if err := format.Node(&newFuncBuf, fset, commentedNode); err != nil {
628647
return nil, nil, err
629648
}
630-
// Find all the comments within the range and print them to be put somewhere.
631-
// TODO(suzmue): print these in the extracted function at the correct place.
632-
for _, cg := range file.Comments {
633-
if cg.Pos().IsValid() && cg.Pos() < end && cg.Pos() >= start {
634-
for _, c := range cg.List {
635-
fmt.Fprintln(&commentBuf, c.Text)
636-
}
637-
}
638-
}
639649

640650
// We're going to replace the whole enclosing function,
641651
// 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
11871197
return isOverriden
11881198
}
11891199

1190-
// parseBlockStmt generates an AST file from the given text. We then return the portion of the
1191-
// file that represents the text.
1192-
func parseBlockStmt(fset *token.FileSet, src []byte) (*ast.BlockStmt, error) {
1200+
// parseStmts parses the specified source (a list of statements) and
1201+
// returns them as a BlockStmt along with any associated comments.
1202+
func parseStmts(fset *token.FileSet, src []byte) (*ast.BlockStmt, []*ast.CommentGroup, error) {
11931203
text := "package main\nfunc _() { " + string(src) + " }"
1194-
extract, err := parser.ParseFile(fset, "", text, parser.SkipObjectResolution)
1204+
file, err := parser.ParseFile(fset, "", text, parser.ParseComments|parser.SkipObjectResolution)
11951205
if err != nil {
1196-
return nil, err
1206+
return nil, nil, err
11971207
}
1198-
if len(extract.Decls) == 0 {
1199-
return nil, fmt.Errorf("parsed file does not contain any declarations")
1208+
if len(file.Decls) != 1 {
1209+
return nil, nil, fmt.Errorf("got %d declarations, want 1", len(file.Decls))
12001210
}
1201-
decl, ok := extract.Decls[0].(*ast.FuncDecl)
1211+
decl, ok := file.Decls[0].(*ast.FuncDecl)
12021212
if !ok {
1203-
return nil, fmt.Errorf("parsed file does not contain expected function declaration")
1213+
return nil, nil, bug.Errorf("parsed file does not contain expected function declaration")
12041214
}
12051215
if decl.Body == nil {
1206-
return nil, fmt.Errorf("extracted function has no body")
1216+
return nil, nil, bug.Errorf("extracted function has no body")
12071217
}
1208-
return decl.Body, nil
1218+
return decl.Body, file.Comments, nil
12091219
}
12101220

12111221
// generateReturnInfo generates the information we need to adjust the return statements and

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

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -237,20 +237,14 @@ func (b *B) LongListWithT(ctx context.Context, t *testing.T) (int, error) {
237237
+}
238238
+
239239
-- @contextFuncB/context.go --
240-
@@ -33 +33,6 @@
241-
- sum := b.x + b.y //@loc(B_AddPWithB, re`(?s:^.*?Err\(\))`)
242-
+ //@loc(B_AddPWithB, re`(?s:^.*?Err\(\))`)
240+
@@ -33 +33,4 @@
243241
+ return newFunction(ctx, tB, b)
244242
+}
245243
+
246244
+func newFunction(ctx context.Context, tB *testing.B, b *B) (int, error) {
247-
+ sum := b.x + b.y
248245
-- @contextFuncT/context.go --
249-
@@ -42 +42,6 @@
250-
- p4 := p1 + p2 //@loc(B_LongListWithT, re`(?s:^.*?Err\(\))`)
251-
+ //@loc(B_LongListWithT, re`(?s:^.*?Err\(\))`)
246+
@@ -42 +42,4 @@
252247
+ return newFunction(ctx, t, p1, p2, p3)
253248
+}
254249
+
255250
+func newFunction(ctx context.Context, t *testing.T, p1 int, p2 int, p3 int) (int, error) {
256-
+ p4 := p1 + p2

0 commit comments

Comments
 (0)