internal/lsp: print comments that would be lost during extract func

Due to the limitations of comments in ast, it is difficult to move
comments. The extract function feature currently does not handle
comments at all. This change instead prints the comments that would
have been lost above the call to the function, so that the user can
easily recover them. Otherwise, it was possible for users to lose
comments and not notice.

Updates golang/go#37170

Change-Id: I1e2d865f5deddefbb0417732490decbdfcde5f3d
Reviewed-on: https://go-review.googlesource.com/c/tools/+/313211
Trust: Suzy Mueller <suzmue@golang.org>
Run-TryBot: Suzy Mueller <suzmue@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>

Author: suzmue

internal/lsp/source/extract.go

@@ -487,7 +487,7 @@ func extractFunction(fset *token.FileSet, rng span.Range, src []byte, file *ast.
 		declarations = initializeVars(uninitialized, retVars, seenUninitialized, seenVars)
 	}
 
-	var declBuf, replaceBuf, newFuncBuf, ifBuf bytes.Buffer
+	var declBuf, replaceBuf, newFuncBuf, ifBuf, commentBuf bytes.Buffer
 	if err := format.Node(&declBuf, fset, declarations); err != nil {
 		return nil, err
 	}
@@ -502,6 +502,15 @@ func extractFunction(fset *token.FileSet, rng span.Range, src []byte, file *ast.
 	if err := format.Node(&newFuncBuf, fset, newFunc); err != nil {
 		return 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() < rng.End && cg.Pos() >= rng.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.
@@ -513,6 +522,10 @@ func extractFunction(fset *token.FileSet, rng span.Range, src []byte, file *ast.
 
 	var fullReplacement strings.Builder
 	fullReplacement.Write(before)
+	if commentBuf.Len() > 0 {
+		comments := strings.ReplaceAll(commentBuf.String(), "\n", newLineIndent)
+		fullReplacement.WriteString(comments)
+	}
 	if declBuf.Len() > 0 { // add any initializations, if needed
 		initializations := strings.ReplaceAll(declBuf.String(), "\n", newLineIndent) +
 			newLineIndent

internal/lsp/testdata/extract/extract_function/extract_args_returns.go.golden

@@ -3,6 +3,7 @@ package extract
 
 func _() {
 	a := 1
+	//@mark(exSt0, "a")
 	a = fn0(a) //@mark(exEn0, "2")
 	//@extractfunc(exSt0, exEn0)
 	b := a * 2 //@mark(exB, "	b")
@@ -24,6 +25,7 @@ func _() {
 	a = 5     //@mark(exSt0, "a")
 	a = a + 2 //@mark(exEn0, "2")
 	//@extractfunc(exSt0, exEn0)
+	//@mark(exB, "	b")
 	fn0(a)  //@mark(exEnd, "4")
 	//@extractfunc(exB, exEnd)
 }

internal/lsp/testdata/extract/extract_function/extract_basic.go.golden

@@ -2,6 +2,7 @@
 package extract
 
 func _() {
+	//@mark(exSt1, "a")
 	fn0() //@mark(exEn1, "4")
 	//@extractfunc(exSt1, exEn1)
 }

internal/lsp/testdata/extract/extract_function/extract_basic_comment.go

@@ -0,0 +1,8 @@
+package extract
+
+func _() {
+	a := /* comment in the middle of a line */ 1 //@mark(exSt18, "a")
+	// Comment on its own line
+	_ = 3 + 4 //@mark(exEn18, "4")
+	//@extractfunc(exSt18, exEn18)
+}

internal/lsp/testdata/extract/extract_function/extract_basic_comment.go.golden

@@ -0,0 +1,17 @@
+-- functionextraction_extract_basic_comment_4_2 --
+package extract
+
+func _() {
+	/* comment in the middle of a line */
+	//@mark(exSt18, "a")
+	// Comment on its own line
+	fn0() //@mark(exEn18, "4")
+	//@extractfunc(exSt18, exEn18)
+}
+
+func fn0() {
+	a := 1
+
+	_ = 3 + 4
+}
+

internal/lsp/testdata/extract/extract_function/extract_issue_44813.go.golden

@@ -4,6 +4,7 @@ package extract
 import "fmt"
 
 func main() {
+	//@mark(exSt9, "x")
 	x := fn0() //@mark(exEn9, "}")
 	//@extractfunc(exSt9, exEn9)
 	fmt.Printf("%x\n", x)

internal/lsp/testdata/extract/extract_function/extract_return_basic.go.golden

@@ -3,6 +3,7 @@ package extract
 
 func _() bool {
 	x := 1
+	//@mark(exSt2, "if")
 	cond0, ret0 := fn0(x)
 	if cond0 {
 		return ret0

internal/lsp/testdata/extract/extract_function/extract_return_basic_nonnested.go.golden

@@ -2,6 +2,7 @@
 package extract
 
 func _() bool {
+	//@mark(exSt13, "x")
 	return fn0() //@mark(exEn13, "false")
 	//@extractfunc(exSt13, exEn13)
 }

internal/lsp/testdata/extract/extract_function/extract_return_complex.go.golden

@@ -6,6 +6,7 @@ import "fmt"
 func _() (int, string, error) {
 	x := 1
 	y := "hello"
+	//@mark(exSt3, "z")
 	z, cond0, ret0, ret1, ret2 := fn0(y, x)
 	if cond0 {
 		return ret0, ret1, ret2

internal/lsp/testdata/extract/extract_function/extract_return_complex_nonnested.go.golden

@@ -6,6 +6,7 @@ import "fmt"
 func _() (int, string, error) {
 	x := 1
 	y := "hello"
+	//@mark(exSt10, "z")
 	return fn0(y, x) //@mark(exEn10, "nil")
 	//@extractfunc(exSt10, exEn10)
 }

internal/lsp/testdata/extract/extract_function/extract_return_func_lit.go.golden

@@ -5,6 +5,7 @@ import "go/ast"
 
 func _() {
 	ast.Inspect(ast.NewIdent("a"), func(n ast.Node) bool {
+		//@mark(exSt4, "if")
 		cond0, ret0 := fn0(n)
 		if cond0 {
 			return ret0

internal/lsp/testdata/extract/extract_function/extract_return_func_lit_nonnested.go.golden

@@ -5,6 +5,7 @@ import "go/ast"
 
 func _() {
 	ast.Inspect(ast.NewIdent("a"), func(n ast.Node) bool {
+		//@mark(exSt11, "if")
 		return fn0(n) //@mark(exEn11, "false")
 	})
 	//@extractfunc(exSt11, exEn11)

internal/lsp/testdata/extract/extract_function/extract_return_init.go.golden

@@ -3,6 +3,7 @@ package extract
 
 func _() string {
 	x := 1
+	//@mark(exSt5, "if")
 	cond0, ret0 := fn0(x)
 	if cond0 {
 		return ret0

internal/lsp/testdata/extract/extract_function/extract_return_init_nonnested.go.golden

@@ -3,6 +3,7 @@ package extract
 
 func _() string {
 	x := 1
+	//@mark(exSt12, "if")
 	return fn0(x) //@mark(exEn12, "\"b\"")
 	//@extractfunc(exSt12, exEn12)
 }

internal/lsp/testdata/extract/extract_function/extract_smart_initialization.go.golden

@@ -3,6 +3,7 @@ package extract
 
 func _() {
 	var a []int
+	//@mark(exSt6, "a")
 	a, b := fn0(a)           //@mark(exEn6, "4")
 	//@extractfunc(exSt6, exEn6)
 	a = append(a, b)

internal/lsp/testdata/extract/extract_function/extract_smart_return.go.golden

@@ -4,6 +4,7 @@ package extract
 func _() {
 	var b []int
 	var a int
+	//@mark(exSt7, "a")
 	b = fn0(a, b) //@mark(exEn7, ")")
 	b[0] = 1
 	//@extractfunc(exSt7, exEn7)

internal/lsp/testdata/extract/extract_function/extract_unnecessary_param.go.golden

@@ -4,6 +4,7 @@ package extract
 func _() {
 	var b []int
 	var a int
+	//@mark(exSt8, "a")
 	a, b = fn0(b) //@mark(exEn8, ")")
 	b[0] = 1
 	if a == 2 {

internal/lsp/testdata/summary.txt.golden

@@ -14,7 +14,7 @@ FormatCount = 6
 ImportCount = 8
 SemanticTokenCount = 3
 SuggestedFixCount = 40
-FunctionExtractionCount = 17
+FunctionExtractionCount = 18
 DefinitionsCount = 95
 TypeDefinitionsCount = 10
 HighlightsCount = 69