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

@@ -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))
 		}
 	})
 }

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)