internal/lsp: simplify and correct fixing imports for CodeAction

The code was introducting syntax errors for some edge cases (example in
regtest/import_test.go), and I found it hard to follow.

The new code passes all the tests. There are new regtests to guarantee
no CodeActions are returned for some cases that vim testing noticed.

Change-Id: Ia09da667f74673c11bfe185e4f76a76c66940105
Reviewed-on: https://go-review.googlesource.com/c/tools/+/233117
Run-TryBot: Peter Weinberger <pjw@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Heschi Kreinick <heschi@google.com>

internal/lsp/fake/editor.go

@@ -700,3 +700,24 @@ func (e *Editor) CodeLens(ctx context.Context, path string) ([]protocol.CodeLens
 	}
 	return lens, nil
 }
+
+// CodeAction executes a codeAction request on the server.
+func (e *Editor) CodeAction(ctx context.Context, path string) ([]protocol.CodeAction, error) {
+	if e.server == nil {
+		return nil, nil
+	}
+	e.mu.Lock()
+	_, ok := e.buffers[path]
+	e.mu.Unlock()
+	if !ok {
+		return nil, fmt.Errorf("buffer %q is not open", path)
+	}
+	params := &protocol.CodeActionParams{
+		TextDocument: e.textDocumentIdentifier(path),
+	}
+	lens, err := e.server.CodeAction(ctx, params)
+	if err != nil {
+		return nil, err
+	}
+	return lens, nil
+}

internal/lsp/regtest/imports_test.go

@@ -0,0 +1,100 @@
+package regtest
+
+import (
+	"testing"
+)
+
+const needs = `
+-- go.mod --
+module foo
+
+-- a.go --
+package main
+func f() {}
+`
+const ntest = `package main
+func TestZ(t *testing.T) {
+	f()
+}
+`
+const want = `package main
+
+import "testing"
+
+func TestZ(t *testing.T) {
+	f()
+}
+`
+
+func TestIssue38815(t *testing.T) {
+	// it was returning
+	// "package main\nimport \"testing\"\npackage main..."
+	runner.Run(t, needs, func(t *testing.T, env *Env) {
+		env.CreateBuffer("a_test.go", ntest)
+		env.SaveBuffer("a_test.go")
+		got := env.Editor.BufferText("a_test.go")
+		if want != got {
+			t.Errorf("got\n%q, wanted\n%q", got, want)
+		}
+	})
+}
+
+const vim1 = `package main
+
+import "fmt"
+
+var foo = 1
+var bar = 2
+
+func main() {
+	fmt.Printf("This is a test %v\n", foo)
+	fmt.Printf("This is another test %v\n", foo)
+	fmt.Printf("This is also a test %v\n", foo)
+}
+`
+
+func TestVim1(t *testing.T) {
+	// The file remains unchanged, but if there are any CodeActions returned, they confuse vim.
+	// Therefore check for no CodeActions
+	runner.Run(t, vim1, func(t *testing.T, env *Env) {
+		env.CreateBuffer("main.go", vim1)
+		env.OrganizeImports("main.go")
+		actions := env.CodeAction("main.go")
+		if len(actions) > 0 {
+			got := env.Editor.BufferText("main.go")
+			t.Errorf("unexpected actions %#v", actions)
+			if got == vim1 {
+				t.Errorf("no changes")
+			} else {
+				t.Errorf("got\n%q", got)
+				t.Errorf("was\n%q", vim1)
+			}
+		}
+	})
+}
+
+const vim2 = `package main
+
+import (
+	"fmt"
+
+	"example.com/blah"
+
+	"rubbish.com/useless"
+)
+
+func main() {
+	fmt.Println(blah.Name, useless.Name)
+}
+`
+
+func TestVim2(t *testing.T) {
+	runner.Run(t, vim1, func(t *testing.T, env *Env) {
+		env.CreateBuffer("main.go", vim2)
+		env.OrganizeImports("main.go")
+		actions := env.CodeAction("main.go")
+		if len(actions) > 0 {
+			t.Errorf("unexpected actions %#v", actions)
+		}
+	})
+}

internal/lsp/regtest/wrappers.go

@@ -187,3 +187,14 @@ func (e *Env) CodeLens(path string) []protocol.CodeLens {
 	}
 	return lens
 }
+
+// CodeAction calls testDocument/codeAction for the given path, and calls
+// t.Fatal if there are errors.
+func (e *Env) CodeAction(path string) []protocol.CodeAction {
+	e.T.Helper()
+	actions, err := e.Editor.CodeAction(e.Ctx, path)
+	if err != nil {
+		e.T.Fatal(err)
+	}
+	return actions
+}

internal/lsp/source/format.go

@@ -11,14 +11,12 @@ import (
 	"go/ast"
 	"go/format"
 	"go/parser"
-	"go/scanner"
 	"go/token"
 
 	"golang.org/x/tools/internal/event"
 	"golang.org/x/tools/internal/imports"
 	"golang.org/x/tools/internal/lsp/diff"
 	"golang.org/x/tools/internal/lsp/protocol"
-	"golang.org/x/tools/internal/span"
 	errors "golang.org/x/xerrors"
 )
 
@@ -100,7 +98,7 @@ func computeImportEdits(ctx context.Context, view View, ph ParseGoHandle, option
 	if err != nil {
 		return nil, nil, err
 	}
-	origAST, _, origMapper, _, err := ph.Parse(ctx)
+	_, _, origMapper, _, err := ph.Parse(ctx)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -110,8 +108,7 @@ func computeImportEdits(ctx context.Context, view View, ph ParseGoHandle, option
 		return nil, nil, err
 	}
 
-	origImports, origImportOffset := trimToImports(view.Session().Cache().FileSet(), origAST, origData)
-	allFixEdits, err = computeFixEdits(view, ph, options, origData, origAST, origMapper, origImports, origImportOffset, allFixes)
+	allFixEdits, err = computeFixEdits(view, ph, options, origData, origMapper, allFixes)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -119,7 +116,7 @@ func computeImportEdits(ctx context.Context, view View, ph ParseGoHandle, option
 	// Apply all of the import fixes to the file.
 	// Add the edits for each fix to the result.
 	for _, fix := range allFixes {
-		edits, err := computeFixEdits(view, ph, options, origData, origAST, origMapper, origImports, origImportOffset, []*imports.ImportFix{fix})
+		edits, err := computeFixEdits(view, ph, options, origData, origMapper, []*imports.ImportFix{fix})
 		if err != nil {
 			return nil, nil, err
 		}
@@ -136,11 +133,10 @@ func computeOneImportFixEdits(ctx context.Context, view View, ph ParseGoHandle,
 	if err != nil {
 		return nil, err
 	}
-	origAST, _, origMapper, _, err := ph.Parse(ctx)
+	_, _, origMapper, _, err := ph.Parse(ctx) // ph.Parse returns values never used
 	if err != nil {
 		return nil, err
 	}
-	origImports, origImportOffset := trimToImports(view.Session().Cache().FileSet(), origAST, origData)
 
 	options := &imports.Options{
 		// Defaults.
@@ -151,154 +147,61 @@ func computeOneImportFixEdits(ctx context.Context, view View, ph ParseGoHandle,
 		TabIndent:  true,
 		TabWidth:   8,
 	}
-	return computeFixEdits(view, ph, options, origData, origAST, origMapper, origImports, origImportOffset, []*imports.ImportFix{fix})
+	return computeFixEdits(view, ph, options, origData, origMapper, []*imports.ImportFix{fix})
 }
 
-func computeFixEdits(view View, ph ParseGoHandle, options *imports.Options, origData []byte, origAST *ast.File, origMapper *protocol.ColumnMapper, origImports []byte, origImportOffset int, fixes []*imports.ImportFix) ([]protocol.TextEdit, error) {
+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 = append(fixedData, '\n') // ApplyFixes comes out missing the newline, go figure.
 	if err != nil {
 		return nil, err
 	}
-	fixedFset := token.NewFileSet()
-	fixedAST, err := parser.ParseFile(fixedFset, filename, fixedData, parser.ImportsOnly)
-	// Any error here prevents us from computing the edits.
-	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.
 	}
-	fixedImports, fixedImportsOffset := trimToImports(fixedFset, fixedAST, fixedData)
-
-	// Prepare the diff. If both sides had import statements, we can diff
-	// just those sections against each other, then shift the resulting
-	// edits to the right lines in the original file.
-	left, right := origImports, fixedImports
-
-	// If there is no diff, return early, as there's no need to compute edits.
-	// imports.ApplyFixes also formats the file, and this way we avoid
-	// unnecessary formatting, which may cause further issues if we can't
-	// find an import block on which to anchor the diffs.
-	if len(left) == 0 && len(right) == 0 {
-		return nil, nil
-	}
-
-	converter := span.NewContentConverter(filename, origImports)
-	offset := origImportOffset
-
-	// If one side or the other has no imports, we won't know where to
-	// anchor the diffs. Instead, use the beginning of the file, up to its
-	// first non-imports decl. We know the imports code will insert
-	// somewhere before that.
-	if origImportOffset == 0 || fixedImportsOffset == 0 {
-		left, _ = trimToFirstNonImport(view.Session().Cache().FileSet(), origAST, origData, nil)
-		fixedData, err = imports.ApplyFixes(fixes, filename, origData, options, 0)
-		if err != nil {
-			return nil, err
-		}
-		// We need the whole file here, not just the ImportsOnly versions we made above.
-		fixedAST, err = parser.ParseFile(fixedFset, filename, fixedData, 0)
-		if fixedAST == nil {
-			return nil, err
-		}
-		var ok bool
-		right, ok = trimToFirstNonImport(fixedFset, fixedAST, fixedData, err)
-		if !ok {
-			return nil, errors.Errorf("error %v detected in the import block", err)
-		}
-		// We're now working with a prefix of the original file, so we can
-		// use the original converter, and there is no offset on the edits.
-		converter = origMapper.Converter
-		offset = 0
-	}
-
-	// Perform the diff and adjust the results for the trimming, if any.
-	edits := view.Options().ComputeEdits(ph.File().Identity().URI, string(left), string(right))
-	for i := range edits {
-		s, err := edits[i].Span.WithPosition(converter)
-		if err != nil {
-			return nil, err
-		}
-		start := span.NewPoint(s.Start().Line()+offset, s.Start().Column(), -1)
-		end := span.NewPoint(s.End().Line()+offset, s.End().Column(), -1)
-		edits[i].Span = span.New(s.URI(), start, end)
+	// trim the original data to match fixedData
+	left := importPrefix(filename, 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))
 	return ToProtocolEdits(origMapper, edits)
 }
 
-// trimToImports returns a section of the source file that covers all of the
-// import declarations, and the line offset into the file that section starts at.
-func trimToImports(fset *token.FileSet, f *ast.File, src []byte) ([]byte, int) {
-	var firstImport, lastImport ast.Decl
-	for _, decl := range f.Decls {
-		if gen, ok := decl.(*ast.GenDecl); ok && gen.Tok == token.IMPORT {
-			if firstImport == nil {
-				firstImport = decl
-			}
-			lastImport = decl
-		}
+// 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 {
+	fset := token.NewFileSet()
+	// do as little parsing as possible
+	f, err := parser.ParseFile(fset, filename, src, parser.ImportsOnly|parser.ParseComments)
+	if err != nil { // This can happen if 'package' is misspelled
+		return ""
 	}
-
-	if firstImport == nil {
-		return nil, 0
-	}
-	tok := fset.File(f.Pos())
-	start := firstImport.Pos()
-	end := lastImport.End()
-	// The parser will happily feed us nonsense. See golang/go#36610.
-	tokStart, tokEnd := token.Pos(tok.Base()), token.Pos(tok.Base()+tok.Size())
-	if start < tokStart || start > tokEnd || end < tokStart || end > tokEnd {
-		return nil, 0
-	}
-	if nextLine := fset.Position(end).Line + 1; tok.LineCount() >= nextLine {
-		end = fset.File(f.Pos()).LineStart(nextLine)
-	}
-	if start > end {
-		return nil, 0
-	}
-
-	startLineOffset := fset.Position(start).Line - 1 // lines are 1-indexed.
-	return src[fset.Position(start).Offset:fset.Position(end).Offset], startLineOffset
-}
-
-// trimToFirstNonImport returns src from the beginning to the first non-import
-// declaration, or the end of the file if there is no such decl.
-func trimToFirstNonImport(fset *token.FileSet, f *ast.File, src []byte, err error) ([]byte, bool) {
-	var firstDecl ast.Decl
-	for _, decl := range f.Decls {
-		if gen, ok := decl.(*ast.GenDecl); ok && gen.Tok == token.IMPORT {
-			continue
-		}
-		firstDecl = decl
-		break
-	}
-	tok := fset.File(f.Pos())
-	if tok == nil {
-		return nil, false
-	}
-	end := f.End()
-	if firstDecl != nil {
-		if firstDeclLine := fset.Position(firstDecl.Pos()).Line; firstDeclLine > 1 {
-			end = tok.LineStart(firstDeclLine - 1)
-		}
-	}
-	// Any errors in the file must be after the part of the file that we care about.
-	switch err := err.(type) {
-	case *scanner.Error:
-		pos := tok.Pos(err.Pos.Offset)
-		if pos <= end {
-			return nil, false
-		}
-	case scanner.ErrorList:
-		if err.Len() > 0 {
-			pos := tok.Pos(err[0].Pos.Offset)
-			if pos <= end {
-				return nil, false
+	myStart := fset.File(f.Pos()).Base() // 1, but the generality costs little
+	pkgEnd := int(f.Name.NamePos) + len(f.Name.Name)
+	var importEnd int
+	for _, d := range f.Decls {
+		if x, ok := d.(*ast.GenDecl); ok && x.Tok == token.IMPORT {
+			e := int(d.End()) - myStart
+			if e > importEnd {
+				importEnd = e
 			}
 		}
 	}
-	return src[0:fset.Position(end).Offset], true
+	if importEnd > 0 {
+		return string(src[:importEnd])
+	}
+	// 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())
+		}
+	}
+	return string(src[:pkgEnd])
 }
 
 func computeTextEdits(ctx context.Context, view View, fh FileHandle, m *protocol.ColumnMapper, formatted string) ([]protocol.TextEdit, error) {

internal/lsp/source/format_test.go

@@ -1,25 +1,28 @@
 package source
 
 import (
-	"go/parser"
-	"go/token"
 	"testing"
 )
 
-func TestTrimToImports(t *testing.T) {
-	const input = `package source
-
-import (
-	m
-	"fmt"
-)
-
-func foo() {
-	fmt.Println("hi")
+type data struct {
+	input, want string
 }
-`
 
-	fs := token.NewFileSet()
-	f, _ := parser.ParseFile(fs, "foo.go", input, parser.ImportsOnly)
-	trimToImports(fs, f, []byte(input))
+func TestImportPrefix(t *testing.T) {
+	var tdata = []data{
+		{"package foo\n", "package foo\n"},
+		{"package foo\n\nfunc f(){}\n", "package foo\n"},
+		{"package foo\n\nimport \"fmt\"\n", "package foo\n\nimport \"fmt\""},
+		{"package foo\nimport (\n\"fmt\"\n)\n", "package foo\nimport (\n\"fmt\"\n)"},
+		{"\n\n\npackage foo\n", "\n\n\npackage foo\n"},
+		{"// 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"},
+	}
+	for i, x := range tdata {
+		got := importPrefix("a.go", []byte(x.input))
+		if got != x.want {
+			t.Errorf("%d: got\n%q, wanted\n%q", i, got, x.want)
+		}
+	}
 }

internal/lsp/testdata/lsp/primarymod/imports/issue35458.go.golden

@@ -1,9 +1,5 @@
 -- goimports --
-
-
-
-
-
+// package doc
 package imports //@import("package")
 
 

internal/lsp/testdata/lsp/primarymod/imports/issue35458.go.in

@@ -3,6 +3,7 @@
 
 
 
+// package doc
 package imports //@import("package")