internal/lsp/source: don't format the whole file when adding imports

We want people to add imports as they need them. That means we probably
don't want adding an import to reformat your whole file while you're in
the middle of editing it.

Unfortunately, the AST package doesn't offer any help with this --
there's no good way to get a diff out of it. Instead, we apply the
changes, then diff a subset of the file. Picking that subset is tricky,
see the code for details.

Also delete a dead function, Imports, which should have been unused but
was still being called in tests.

Fixes golang/go#30843.

Change-Id: I09a5344e910f65510003c4006ea5b11657922315
Reviewed-on: https://go-review.googlesource.com/c/tools/+/205678
Reviewed-by: Rebecca Stambler <[email protected]>

internal/lsp/source/format.go

@@ -8,7 +8,10 @@ package source
 import (
 	"bytes"
 	"context"
+	"go/ast"
 	"go/format"
+	"go/parser"
+	"go/token"
 
 	"golang.org/x/tools/internal/imports"
 	"golang.org/x/tools/internal/lsp/diff"
@@ -84,35 +87,44 @@ func formatSource(ctx context.Context, s Snapshot, f File) ([]byte, error) {
 	return format.Source(data)
 }
 
-// Imports formats a file using the goimports tool.
-func Imports(ctx context.Context, view View, f File) ([]protocol.TextEdit, error) {
-	ctx, done := trace.StartSpan(ctx, "source.Imports")
+type ImportFix struct {
+	Fix   *imports.ImportFix
+	Edits []protocol.TextEdit
+}
+
+// AllImportsFixes formats f for each possible fix to the imports.
+// In addition to returning the result of applying all edits,
+// it returns a list of fixes that could be applied to the file, with the
+// corresponding TextEdits that would be needed to apply that fix.
+func AllImportsFixes(ctx context.Context, view View, f File) (allFixEdits []protocol.TextEdit, editsPerFix []*ImportFix, err error) {
+	ctx, done := trace.StartSpan(ctx, "source.AllImportsFixes")
 	defer done()
 
 	_, cphs, err := view.CheckPackageHandles(ctx, f)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	cph, err := NarrowestCheckPackageHandle(cphs)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	pkg, err := cph.Check(ctx)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	if hasListErrors(pkg) {
-		return nil, errors.Errorf("%s has list errors, not running goimports", f.URI())
+		return nil, nil, errors.Errorf("%s has list errors, not running goimports", f.URI())
 	}
-	ph, err := pkg.File(f.URI())
-	if err != nil {
-		return nil, err
+	var ph ParseGoHandle
+	for _, h := range pkg.Files() {
+		if h.File().Identity().URI == f.URI() {
+			ph = h
+		}
 	}
-	// Be extra careful that the file's ParseMode is correct,
-	// otherwise we might replace the user's code with a trimmed AST.
-	if ph.Mode() != ParseFull {
-		return nil, errors.Errorf("%s was parsed in the incorrect mode", ph.File().Identity().URI)
+	if ph == nil {
+		return nil, nil, errors.Errorf("no ParseGoHandle for %s", f.URI())
 	}
+
 	options := &imports.Options{
 		// Defaults.
 		AllErrors:  true,
@@ -122,122 +134,150 @@ func Imports(ctx context.Context, view View, f File) ([]protocol.TextEdit, error
 		TabIndent:  true,
 		TabWidth:   8,
 	}
-	var formatted []byte
-	importFn := func(opts *imports.Options) error {
-		data, _, err := ph.File().Read(ctx)
-		if err != nil {
-			return err
-		}
-		formatted, err = imports.Process(ph.File().Identity().URI.Filename(), data, opts)
+	err = view.RunProcessEnvFunc(ctx, func(opts *imports.Options) error {
+		allFixEdits, editsPerFix, err = computeImportEdits(ctx, view, ph, opts)
 		return err
-	}
-	err = view.RunProcessEnvFunc(ctx, importFn, options)
-	if err != nil {
-		return nil, err
-	}
-	_, m, _, err := ph.Parse(ctx)
+	}, options)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
-	return computeTextEdits(ctx, view, ph.File(), m, string(formatted))
-}
 
-type ImportFix struct {
-	Fix   *imports.ImportFix
-	Edits []protocol.TextEdit
+	return allFixEdits, editsPerFix, nil
 }
 
-// AllImportsFixes formats f for each possible fix to the imports.
-// In addition to returning the result of applying all edits,
-// it returns a list of fixes that could be applied to the file, with the
-// corresponding TextEdits that would be needed to apply that fix.
-func AllImportsFixes(ctx context.Context, view View, f File) (edits []protocol.TextEdit, editsPerFix []*ImportFix, err error) {
-	ctx, done := trace.StartSpan(ctx, "source.AllImportsFixes")
-	defer done()
+// computeImportEdits computes a set of edits that perform one or all of the
+// necessary import fixes.
+func computeImportEdits(ctx context.Context, view View, ph ParseGoHandle, options *imports.Options) (allFixEdits []protocol.TextEdit, editsPerFix []*ImportFix, err error) {
+	filename := ph.File().Identity().URI.Filename()
 
-	_, cphs, err := view.CheckPackageHandles(ctx, f)
+	// Build up basic information about the original file.
+	origData, _, err := ph.File().Read(ctx)
 	if err != nil {
 		return nil, nil, err
 	}
-	cph, err := NarrowestCheckPackageHandle(cphs)
-	if err != nil {
-		return nil, nil, err
-	}
-	pkg, err := cph.Check(ctx)
+	origAST, origMapper, _, err := ph.Parse(ctx)
 	if err != nil {
 		return nil, nil, err
 	}
-	if hasListErrors(pkg) {
-		return nil, nil, errors.Errorf("%s has list errors, not running goimports", f.URI())
-	}
-	options := &imports.Options{
-		// Defaults.
-		AllErrors:  true,
-		Comments:   true,
-		Fragment:   true,
-		FormatOnly: false,
-		TabIndent:  true,
-		TabWidth:   8,
-	}
-	importFn := func(opts *imports.Options) error {
-		var ph ParseGoHandle
-		for _, h := range pkg.Files() {
-			if h.File().Identity().URI == f.URI() {
-				ph = h
-			}
-		}
-		if ph == nil {
-			return errors.Errorf("no ParseGoHandle for %s", f.URI())
-		}
-		data, _, err := ph.File().Read(ctx)
-		if err != nil {
-			return err
-		}
-		fixes, err := imports.FixImports(f.URI().Filename(), data, opts)
-		if err != nil {
-			return err
-		}
-		// Do not change the file if there are no import fixes.
-		if len(fixes) == 0 {
-			return nil
-		}
-		// Apply all of the import fixes to the file.
-		formatted, err := imports.ApplyFixes(fixes, f.URI().Filename(), data, options)
-		if err != nil {
-			return err
-		}
-		_, m, _, err := ph.Parse(ctx)
+	origImports, origImportOffset := trimToImports(view.Session().Cache().FileSet(), origAST, origData)
+
+	computeFixEdits := func(fixes []*imports.ImportFix) ([]protocol.TextEdit, error) {
+		// Apply the fixes and re-parse the file so that we can locate the
+		// new imports.
+		fixedData, err := imports.ApplyFixes(fixes, filename, origData, options)
 		if err != nil {
-			return err
+			return nil, err
 		}
-		edits, err = computeTextEdits(ctx, view, ph.File(), m, string(formatted))
+		fixedFset := token.NewFileSet()
+		fixedAST, err := parser.ParseFile(fixedFset, filename, fixedData, parser.ImportsOnly)
 		if err != nil {
-			return err
+			return nil, err
 		}
-		// Add the edits for each fix to the result.
-		editsPerFix = make([]*ImportFix, len(fixes))
-		for i, fix := range fixes {
-			formatted, err := imports.ApplyFixes([]*imports.ImportFix{fix}, f.URI().Filename(), data, options)
+		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
+		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)
+			// We need the whole AST here, not just the ImportsOnly AST we parsed above.
+			fixedAST, err = parser.ParseFile(fixedFset, filename, fixedData, 0)
 			if err != nil {
-				return err
+				return nil, err
 			}
-			edits, err := computeTextEdits(ctx, view, ph.File(), m, string(formatted))
+			right = trimToFirstNonImport(fixedFset, fixedAST, fixedData)
+			// 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 err
-			}
-			editsPerFix[i] = &ImportFix{
-				Fix:   fix,
-				Edits: edits,
+				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)
 		}
-		return nil
+		return ToProtocolEdits(origMapper, edits)
 	}
-	err = view.RunProcessEnvFunc(ctx, importFn, options)
+
+	allFixes, err := imports.FixImports(filename, origData, options)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	return edits, editsPerFix, nil
+	allFixEdits, err = computeFixEdits(allFixes)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	// 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([]*imports.ImportFix{fix})
+		if err != nil {
+			return nil, nil, err
+		}
+		editsPerFix = append(editsPerFix, &ImportFix{
+			Fix:   fix,
+			Edits: edits,
+		})
+	}
+	return allFixEdits, editsPerFix, nil
+}
+
+// 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
+		}
+	}
+
+	if firstImport == nil {
+		return nil, 0
+	}
+	start := firstImport.Pos()
+	end := fset.File(f.Pos()).LineStart(fset.Position(lastImport.End()).Line + 1)
+	startLineOffset := fset.Position(start).Line - 1 // lines are 1-indexed.
+	return src[fset.Position(firstImport.Pos()).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) []byte {
+	var firstDecl ast.Decl
+	for _, decl := range f.Decls {
+		if gen, ok := decl.(*ast.GenDecl); ok && gen.Tok == token.IMPORT {
+			continue
+		}
+		firstDecl = decl
+		break
+	}
+
+	end := f.End()
+	if firstDecl != nil {
+		end = fset.File(f.Pos()).LineStart(fset.Position(firstDecl.Pos()).Line - 1)
+	}
+	return src[fset.Position(f.Pos()).Offset:fset.Position(end).Offset]
 }
 
 // CandidateImports returns every import that could be added to filename.

internal/lsp/source/source_test.go

@@ -450,22 +450,14 @@ func (r *runner) Import(t *testing.T, spn span.Span) {
 	ctx := r.ctx
 	uri := spn.URI()
 	filename := uri.Filename()
-	goimported := string(r.data.Golden("goimports", filename, func() ([]byte, error) {
-		cmd := exec.Command("goimports", filename)
-		out, _ := cmd.Output() // ignore error, sometimes we have intentionally ungofmt-able files
-		return out, nil
-	}))
 	f, err := r.view.GetFile(ctx, uri)
 	if err != nil {
 		t.Fatalf("failed for %v: %v", spn, err)
 	}
 	fh := r.view.Snapshot().Handle(r.ctx, f)
-	edits, err := source.Imports(ctx, r.view, f)
+	edits, _, err := source.AllImportsFixes(ctx, r.view, f)
 	if err != nil {
-		if goimported != "" {
-			t.Error(err)
-		}
-		return
+		t.Error(err)
 	}
 	data, _, err := fh.Read(ctx)
 	if err != nil {
@@ -480,8 +472,11 @@ func (r *runner) Import(t *testing.T, spn span.Span) {
 		t.Error(err)
 	}
 	got := diff.ApplyEdits(string(data), diffEdits)
-	if goimported != got {
-		t.Errorf("import failed for %s, expected:\n%v\ngot:\n%v", filename, goimported, got)
+	want := string(r.data.Golden("goimports", filename, func() ([]byte, error) {
+		return []byte(got), nil
+	}))
+	if want != got {
+		t.Errorf("import failed for %s, expected:\n%v\ngot:\n%v", filename, want, got)
 	}
 }
 

internal/lsp/testdata/imports/add_import.go.golden

@@ -0,0 +1,13 @@
+-- goimports --
+package imports //@import("package")
+
+import (
+	"bytes"
+	"fmt"
+)
+
+func _() {
+	fmt.Println("")
+	bytes.NewBuffer(nil)
+}
+

internal/lsp/testdata/imports/add_import.go.in

@@ -0,0 +1,10 @@
+package imports //@import("package")
+
+import (
+	"fmt"
+)
+
+func _() {
+	fmt.Println("")
+	bytes.NewBuffer(nil)
+}

internal/lsp/testdata/imports/good_imports.go.golden

@@ -4,6 +4,6 @@ package imports //@import("package")
 import "fmt"
 
 func _() {
-	fmt.Println("")
+fmt.Println("")
 }
 

internal/lsp/testdata/imports/good_imports.go.in

@@ -0,0 +1,7 @@
+package imports //@import("package")
+
+import "fmt"
+
+func _() {
+fmt.Println("")
+}

internal/lsp/testdata/imports/multiple_blocks.go.golden

@@ -0,0 +1,9 @@
+-- goimports --
+package imports //@import("package")
+
+import "fmt"
+
+func _() {
+	fmt.Println("")
+}
+