From 253fce384c97181fd4207b7214f9aaa3aa44ec9b Mon Sep 17 00:00:00 2001
From: pjw <pjw@google.com>
Date: Sun, 24 May 2020 17:13:47 -0400
Subject: [PATCH] internal/lsp: fix formatting edge cases (36824)

One line legal code like `package x; import "os"; func f() {}` was
being misformatted. In these cases the parse flag ImportsOnly loses
important parts of the code, while full parsing works. Presumably
all these cases are short enough that there is no appreciable penalty
from the extra parsing.

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

Change-Id: I9a4581d67c590578f8fdea5ed2a2a58e0bc3c40b
Reviewed-on: https://go-review.googlesource.com/c/tools/+/234900
Run-TryBot: Peter Weinberger <pjw@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Heschi Kreinick <heschi@google.com>
---
 internal/lsp/regtest/formatting_test.go | 24 ++++++++++++++++++++++--
 internal/lsp/source/format.go           | 22 ++++++++++++++++------
 2 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/internal/lsp/regtest/formatting_test.go b/internal/lsp/regtest/formatting_test.go
index 18b3da06e5..4f200fc2f2 100644
--- a/internal/lsp/regtest/formatting_test.go
+++ b/internal/lsp/regtest/formatting_test.go
@@ -60,7 +60,6 @@ func f() {}
 
 // Tests golang/go#36824.
 func TestFormattingOneLineImports36824(t *testing.T) {
-	t.Skipf("golang/go#36824 has not been fixed yet")
 
 	const onelineProgramA = `
 -- a.go --
@@ -79,7 +78,28 @@ func f() { fmt.Println() }
 		got := env.Editor.BufferText("a.go")
 		want := env.ReadWorkspaceFile("a.go.imported")
 		if got != want {
-			t.Errorf("OneLineImports:\n%s", tests.Diff(want, got))
+			t.Errorf("OneLineImports3824:\n%s", tests.Diff(want, got))
+		}
+	})
+}
+
+func TestFormattingOneLineRmImports36824(t *testing.T) {
+	const onelineProgramB = `
+-- a.go --
+package x; import "os"; func f() {}
+
+-- a.go.imported --
+package x
+
+func f() {}
+`
+	runner.Run(t, onelineProgramB, func(t *testing.T, env *Env) {
+		env.OpenFile("a.go")
+		env.OrganizeImports("a.go")
+		got := env.Editor.BufferText("a.go")
+		want := env.ReadWorkspaceFile("a.go.imported")
+		if got != want {
+			t.Errorf("OneLineRmImports:\n%s", tests.Diff(want, got))
 		}
 	})
 }
diff --git a/internal/lsp/source/format.go b/internal/lsp/source/format.go
index 7b89a8e94e..51814a6e2e 100644
--- a/internal/lsp/source/format.go
+++ b/internal/lsp/source/format.go
@@ -12,6 +12,7 @@ import (
 	"go/format"
 	"go/parser"
 	"go/token"
+	"strings"
 
 	"golang.org/x/tools/internal/event"
 	"golang.org/x/tools/internal/imports"
@@ -151,20 +152,29 @@ 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) {
+	// trim the original data to match fixedData
+	left := importPrefix(origData)
+	extra := strings.Index(left, "\n") == -1 // one line may have more than imports
+	if extra {
+		left = string(origData)
+	}
+	if len(left) > 0 && left[len(left)-1] != '\n' {
+		left += "\n"
+	}
 	// Apply the fixes and re-parse the file so that we can locate the
 	// new imports.
-	fixedData, err := imports.ApplyFixes(fixes, "", origData, options, parser.ImportsOnly)
+	flags := parser.ImportsOnly
+	if extra {
+		// used all of origData above, use all of it here too
+		flags = 0
+	}
+	fixedData, err := imports.ApplyFixes(fixes, "", origData, options, flags)
 	if err != nil {
 		return nil, err
 	}
 	if fixedData == nil || fixedData[len(fixedData)-1] != '\n' {
 		fixedData = append(fixedData, '\n') // ApplyFixes may miss the newline, go figure.
 	}
-	// trim the original data to match fixedData
-	left := importPrefix(origData)
-	if len(left) > 0 && left[len(left)-1] != '\n' {
-		left += "\n"
-	}
 	uri := ph.File().Identity().URI
 	edits := view.Options().ComputeEdits(uri, left, string(fixedData))
 	return ToProtocolEdits(origMapper, edits)