internal/lsp/source: fix comment update during rename for short variable declarations

*ast.AssignStmt doesn't have an associated comment group. So, we should
try to find and return a comment just before the identifier.

Fixes golang/go#42134

Change-Id: Ie40717a4973ccfdbd99c3df891c2cfffbb21742d
GitHub-Last-Rev: da75fde2dbf3613f3325dbc5930dfc84ea813b90
GitHub-Pull-Request: golang/tools#323
Reviewed-on: https://go-review.googlesource.com/c/tools/+/327229
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
Trust: Rebecca Stambler <rstambler@golang.org>
Trust: Robert Findley <rfindley@google.com>
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>

internal/lsp/source/rename.go

@@ -289,6 +289,38 @@ func (r *renamer) docComment(pkg Package, id *ast.Ident) *ast.CommentGroup {
 				return decl.Doc
 			}
 		case *ast.Ident:
+		case *ast.AssignStmt:
+			// *ast.AssignStmt doesn't have an associated comment group.
+			// So, we try to find a comment just before the identifier.
+
+			// Try to find a comment group only for short variable declarations (:=).
+			if decl.Tok != token.DEFINE {
+				return nil
+			}
+
+			var file *ast.File
+			for _, f := range pkg.GetSyntax() {
+				if f.Pos() <= id.Pos() && id.Pos() <= f.End() {
+					file = f
+					break
+				}
+			}
+			if file == nil {
+				return nil
+			}
+
+			identLine := r.fset.Position(id.Pos()).Line
+			for _, comment := range file.Comments {
+				if comment.Pos() > id.Pos() {
+					// Comment is after the identifier.
+					continue
+				}
+
+				lastCommentLine := r.fset.Position(comment.End()).Line
+				if lastCommentLine+1 == identLine {
+					return comment
+				}
+			}
 		default:
 			return nil
 		}

internal/lsp/testdata/rename/issue42134/1.go

@@ -0,0 +1,8 @@
+package issue42134
+
+func _() {
+	// foo computes things.
+	foo := func() {}
+
+	foo() //@rename("foo", "bar")
+}

internal/lsp/testdata/rename/issue42134/1.go.golden

@@ -0,0 +1,10 @@
+-- bar-rename --
+package issue42134
+
+func _() {
+	// bar computes things.
+	bar := func() {}
+
+	bar() //@rename("foo", "bar")
+}
+

internal/lsp/testdata/rename/issue42134/2.go

@@ -0,0 +1,12 @@
+package issue42134
+
+import "fmt"
+
+func _() {
+	// minNumber is a min number.
+	// Second line.
+	minNumber := min(1, 2)
+	fmt.Println(minNumber) //@rename("minNumber", "res")
+}
+
+func min(a, b int) int { return a }

internal/lsp/testdata/rename/issue42134/2.go.golden

@@ -0,0 +1,14 @@
+-- res-rename --
+package issue42134
+
+import "fmt"
+
+func _() {
+	// res is a min number.
+	// Second line.
+	res := min(1, 2)
+	fmt.Println(res) //@rename("minNumber", "res")
+}
+
+func min(a, b int) int { return a }
+

internal/lsp/testdata/rename/issue42134/3.go

@@ -0,0 +1,11 @@
+package issue42134
+
+func _() {
+	/*
+	tests contains test cases
+	*/
+	tests := []struct { //@rename("tests", "testCases")
+		in, out string
+	}{}
+	_ = tests
+}

internal/lsp/testdata/rename/issue42134/3.go.golden

@@ -0,0 +1,13 @@
+-- testCases-rename --
+package issue42134
+
+func _() {
+	/*
+	testCases contains test cases
+	*/
+	testCases := []struct { //@rename("tests", "testCases")
+		in, out string
+	}{}
+	_ = testCases
+}
+

internal/lsp/testdata/rename/issue42134/4.go

@@ -0,0 +1,8 @@
+package issue42134
+
+func _() {
+	// a is equal to 5. Comment must stay the same
+
+	a := 5
+	_ = a //@rename("a", "b")
+}

internal/lsp/testdata/rename/issue42134/4.go.golden

@@ -0,0 +1,10 @@
+-- b-rename --
+package issue42134
+
+func _() {
+	// a is equal to 5. Comment must stay the same
+
+	b := 5
+	_ = b //@rename("a", "b")
+}
+

internal/lsp/testdata/summary.txt.golden

@@ -19,7 +19,7 @@ DefinitionsCount = 95
 TypeDefinitionsCount = 18
 HighlightsCount = 69
 ReferencesCount = 25
-RenamesCount = 33
+RenamesCount = 37
 PrepareRenamesCount = 7
 SymbolsCount = 5
 WorkspaceSymbolsCount = 20