From 57a9e4404bf7b38f22bbca9af3ddf0dee8e76a04 Mon Sep 17 00:00:00 2001
From: pjw <pjw@google.com>
Date: Tue, 19 May 2020 15:51:07 -0400
Subject: [PATCH] internal/lsp: fix new bug duplicating comments after includes

Fixes https://github.com/golang/go/issues/39147

Change-Id: I6f78efccbabf21dbb00e56a49d88e26ff4733fba
Reviewed-on: https://go-review.googlesource.com/c/tools/+/234584
Reviewed-by: Heschi Kreinick <heschi@google.com>
---
 internal/lsp/source/format.go      | 25 +++++++++++--------------
 internal/lsp/source/format_test.go |  4 +++-
 2 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/internal/lsp/source/format.go b/internal/lsp/source/format.go
index 94ae2fc565..7b89a8e94e 100644
--- a/internal/lsp/source/format.go
+++ b/internal/lsp/source/format.go
@@ -151,10 +151,9 @@ func computeOneImportFixEdits(ctx context.Context, view View, ph ParseGoHandle,
 }
 
 func computeFixEdits(view View, ph ParseGoHandle, options *imports.Options, origData []byte, origMapper *protocol.ColumnMapper, fixes []*imports.ImportFix) ([]protocol.TextEdit, error) {
-	filename := ph.File().Identity().URI.Filename()
 	// Apply the fixes and re-parse the file so that we can locate the
 	// new imports.
-	fixedData, err := imports.ApplyFixes(fixes, filename, origData, options, parser.ImportsOnly)
+	fixedData, err := imports.ApplyFixes(fixes, "", origData, options, parser.ImportsOnly)
 	if err != nil {
 		return nil, err
 	}
@@ -162,21 +161,21 @@ func computeFixEdits(view View, ph ParseGoHandle, options *imports.Options, orig
 		fixedData = append(fixedData, '\n') // ApplyFixes may miss the newline, go figure.
 	}
 	// trim the original data to match fixedData
-	left := importPrefix(filename, origData)
+	left := importPrefix(origData)
 	if len(left) > 0 && left[len(left)-1] != '\n' {
 		left += "\n"
 	}
-	f := view.Options().ComputeEdits
-	edits := f(ph.File().Identity().URI, left, string(fixedData))
+	uri := ph.File().Identity().URI
+	edits := view.Options().ComputeEdits(uri, left, string(fixedData))
 	return ToProtocolEdits(origMapper, edits)
 }
 
 // return the prefix of the src through the last imports, or if there are
 // no imports, through the package statement (and a subsequent comment group)
-func importPrefix(filename string, src []byte) string {
+func importPrefix(src []byte) string {
 	fset := token.NewFileSet()
 	// do as little parsing as possible
-	f, err := parser.ParseFile(fset, filename, src, parser.ImportsOnly|parser.ParseComments)
+	f, err := parser.ParseFile(fset, "", src, parser.ImportsOnly|parser.ParseComments)
 	if err != nil { // This can happen if 'package' is misspelled
 		return ""
 	}
@@ -191,17 +190,15 @@ func importPrefix(filename string, src []byte) string {
 			}
 		}
 	}
-	if importEnd > 0 {
-		return string(src[:importEnd])
+	if importEnd == 0 {
+		importEnd = pkgEnd
 	}
-	// If there are no imports there may still be comments after the package
-	// name. There could be one group before the package statement, and one after.
 	for _, c := range f.Comments {
-		if int(c.End()) > pkgEnd {
-			pkgEnd = int(c.End())
+		if int(c.End()) > importEnd {
+			importEnd = int(c.End())
 		}
 	}
-	return string(src[:pkgEnd])
+	return string(src[:importEnd])
 }
 
 func computeTextEdits(ctx context.Context, view View, fh FileHandle, m *protocol.ColumnMapper, formatted string) ([]protocol.TextEdit, error) {
diff --git a/internal/lsp/source/format_test.go b/internal/lsp/source/format_test.go
index 0d767fd886..845c87c744 100644
--- a/internal/lsp/source/format_test.go
+++ b/internal/lsp/source/format_test.go
@@ -18,9 +18,11 @@ func TestImportPrefix(t *testing.T) {
 		{"// hi \n\npackage foo //xx\nfunc _(){}\n", "// hi \n\npackage foo //xx\n"},
 		{"package foo //hi\n", "package foo //hi\n"},
 		{"//hi\npackage foo\n//a\n\n//b\n", "//hi\npackage foo\n//a\n\n//b\n"},
+		{"package a\n\nimport (\n  \"fmt\"\n)\n//hi\n",
+			"package a\n\nimport (\n  \"fmt\"\n)\n//hi\n"},
 	}
 	for i, x := range tdata {
-		got := importPrefix("a.go", []byte(x.input))
+		got := importPrefix([]byte(x.input))
 		if got != x.want {
 			t.Errorf("%d: got\n%q, wanted\n%q", i, got, x.want)
 		}