@@ -10,6 +10,7 @@ import (
10
10
"go/ast"
11
11
"go/format"
12
12
"go/parser"
13
+ "go/printer"
13
14
"go/token"
14
15
"go/types"
15
16
"slices"
@@ -449,7 +450,8 @@ func extractFunctionMethod(fset *token.FileSet, start, end token.Pos, src []byte
449
450
return nil , nil , err
450
451
}
451
452
selection := src [startOffset :endOffset ]
452
- extractedBlock , err := parseBlockStmt (fset , selection )
453
+
454
+ extractedBlock , extractedComments , err := parseStmts (fset , selection )
453
455
if err != nil {
454
456
return nil , nil , err
455
457
}
@@ -570,26 +572,53 @@ func extractFunctionMethod(fset *token.FileSet, start, end token.Pos, src []byte
570
572
if canDefine {
571
573
sym = token .DEFINE
572
574
}
573
- var name , funName string
575
+ var funName string
574
576
if isMethod {
575
- name = "newMethod"
576
577
// TODO(suzmue): generate a name that does not conflict for "newMethod".
577
- funName = name
578
+ funName = "newMethod"
578
579
} else {
579
- name = "newFunction"
580
- funName , _ = generateAvailableName (start , path , pkg , info , name , 0 )
580
+ funName , _ = generateAvailableName (start , path , pkg , info , "newFunction" , 0 )
581
581
}
582
582
extractedFunCall := generateFuncCall (hasNonNestedReturn , hasReturnValues , params ,
583
583
append (returns , getNames (retVars )... ), funName , sym , receiverName )
584
584
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.
586
615
newFunc := & ast.FuncDecl {
587
616
Name : ast .NewIdent (funName ),
588
617
Type : & ast.FuncType {
589
618
Params : & ast.FieldList {List : paramTypes },
590
619
Results : & ast.FieldList {List : append (returnTypes , getDecls (retVars )... )},
591
620
},
592
- Body : extractedBlock ,
621
+ // Body handled separately -- see above.
593
622
}
594
623
if isMethod {
595
624
var names []* ast.Ident
@@ -603,39 +632,20 @@ func extractFunctionMethod(fset *token.FileSet, start, end token.Pos, src []byte
603
632
}},
604
633
}
605
634
}
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 {
617
636
return nil , nil , err
618
637
}
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 {
620
640
return nil , nil , err
621
641
}
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 ,
626
645
}
627
- if err := format .Node (& newFuncBuf , fset , newFunc ); err != nil {
646
+ if err := format .Node (& newFuncBuf , fset , commentedNode ); err != nil {
628
647
return nil , nil , err
629
648
}
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
- }
639
649
640
650
// We're going to replace the whole enclosing function,
641
651
// 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
1187
1197
return isOverriden
1188
1198
}
1189
1199
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 ) {
1193
1203
text := "package main\n func _() { " + string (src ) + " }"
1194
- extract , err := parser .ParseFile (fset , "" , text , parser .SkipObjectResolution )
1204
+ file , err := parser .ParseFile (fset , "" , text , parser . ParseComments | parser .SkipObjectResolution )
1195
1205
if err != nil {
1196
- return nil , err
1206
+ return nil , nil , err
1197
1207
}
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 ) )
1200
1210
}
1201
- decl , ok := extract .Decls [0 ].(* ast.FuncDecl )
1211
+ decl , ok := file .Decls [0 ].(* ast.FuncDecl )
1202
1212
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" )
1204
1214
}
1205
1215
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" )
1207
1217
}
1208
- return decl .Body , nil
1218
+ return decl .Body , file . Comments , nil
1209
1219
}
1210
1220
1211
1221
// generateReturnInfo generates the information we need to adjust the return statements and
0 commit comments