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 <[email protected]> Run-TryBot: Suzy Mueller <[email protected]> gopls-CI: kokoro <[email protected]> TryBot-Result: Go Bot <[email protected]> Reviewed-by: Rebecca Stambler <[email protected]>
golang · Apr 28, 2021 · 16b25d2 · 16b25d2
1 parent 7c72a84
commit 16b25d2
Show file tree

Hide file tree

Showing 18 changed files with 55 additions and 2 deletions.
diff --git a/internal/lsp/source/extract.go b/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

diff --git a/internal/lsp/testdata/extract/extract_function/extract_args_returns.go.golden b/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)
 }

diff --git a/internal/lsp/testdata/extract/extract_function/extract_basic.go.golden b/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)
 }

diff --git a/internal/lsp/testdata/extract/extract_function/extract_basic_comment.go b/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)
+}
diff --git a/internal/lsp/testdata/extract/extract_function/extract_basic_comment.go.golden b/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
+}
+
diff --git a/internal/lsp/testdata/extract/extract_function/extract_issue_44813.go.golden b/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)

diff --git a/internal/lsp/testdata/extract/extract_function/extract_return_basic.go.golden b/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

diff --git a/internal/lsp/testdata/extract/extract_function/extract_return_basic_nonnested.go.golden b/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)
 }

diff --git a/internal/lsp/testdata/extract/extract_function/extract_return_complex.go.golden b/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

diff --git a/internal/lsp/testdata/extract/extract_function/extract_return_complex_nonnested.go.golden b/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)
 }

diff --git a/internal/lsp/testdata/extract/extract_function/extract_return_func_lit.go.golden b/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

diff --git a/internal/lsp/testdata/extract/extract_function/extract_return_func_lit_nonnested.go.golden b/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)

diff --git a/internal/lsp/testdata/extract/extract_function/extract_return_init.go.golden b/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

diff --git a/internal/lsp/testdata/extract/extract_function/extract_return_init_nonnested.go.golden b/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)
 }

diff --git a/internal/lsp/testdata/extract/extract_function/extract_smart_initialization.go.golden b/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)

diff --git a/internal/lsp/testdata/extract/extract_function/extract_smart_return.go.golden 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)

diff --git a/internal/lsp/testdata/extract/extract_function/extract_unnecessary_param.go.golden b/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 {

diff --git a/internal/lsp/testdata/summary.txt.golden b/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