internal/lsp/regtest: implement formatting and organizeImports

Add two new fake editor commands: Formatting and OrganizeImports, which
delegate to textDocument/formatting and textDocument/codeAction
respectively. Use this in simple regtests, as well as on save.

Implementing this required fixing a broken assumption about text edits
in the editor: previously these edits were incrementally mutating the
buffer, but the correct implementation should simultaneously mutate the
buffer (i.e., all positions in an edit set refer to the starting buffer
state). This never mattered before because we were only operating on one
edit at a time.

Updates golang/go#36879

Change-Id: I6dec343c4e202288fa20c26df2fbafe9340a1bce
Reviewed-on: https://go-review.googlesource.com/c/tools/+/221539
Run-TryBot: Robert Findley <rfindley@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rohan Challa <rohan@golang.org>

internal/lsp/fake/client.go

@@ -98,10 +98,7 @@ func (c *Client) ApplyEdit(ctx context.Context, params *protocol.ApplyWorkspaceE
 	}
 	for _, change := range params.Edit.DocumentChanges {
 		path := c.ws.URIToPath(change.TextDocument.URI)
-		var edits []Edit
-		for _, lspEdit := range change.Edits {
-			edits = append(edits, fromProtocolTextEdit(lspEdit))
-		}
+		edits := convertEdits(change.Edits)
 		c.EditBuffer(ctx, path, edits)
 	}
 	return &protocol.ApplyWorkspaceEditResponse{Applied: true}, nil

internal/lsp/fake/edit.go

@@ -6,6 +6,7 @@ package fake
 
 import (
 	"fmt"
+	"sort"
 	"strings"
 
 	"golang.org/x/tools/internal/lsp/protocol"
@@ -71,7 +72,7 @@ func inText(p Pos, content []string) bool {
 	}
 	// Note the strict right bound: the column indexes character _separators_,
 	// not characters.
-	if p.Column < 0 || p.Column > len(content[p.Line]) {
+	if p.Column < 0 || p.Column > len([]rune(content[p.Line])) {
 		return false
 	}
 	return true
@@ -80,7 +81,21 @@ func inText(p Pos, content []string) bool {
 // editContent implements a simplistic, inefficient algorithm for applying text
 // edits to our buffer representation. It returns an error if the edit is
 // invalid for the current content.
-func editContent(content []string, edit Edit) ([]string, error) {
+func editContent(content []string, edits []Edit) ([]string, error) {
+	newEdits := make([]Edit, len(edits))
+	copy(newEdits, edits)
+	sort.Slice(newEdits, func(i, j int) bool {
+		if newEdits[i].Start.Line < newEdits[j].Start.Line {
+			return true
+		}
+		if newEdits[i].Start.Line > newEdits[j].Start.Line {
+			return false
+		}
+		return newEdits[i].Start.Column < newEdits[j].Start.Column
+	})
+
+	// Validate edits.
+	for _, edit := range newEdits {
 		if edit.End.Line < edit.Start.Line || (edit.End.Line == edit.Start.Line && edit.End.Column < edit.Start.Column) {
 			return nil, fmt.Errorf("invalid edit: end %v before start %v", edit.End, edit.Start)
 		}
@@ -90,10 +105,26 @@ func editContent(content []string, edit Edit) ([]string, error) {
 		if !inText(edit.End, content) {
 			return nil, fmt.Errorf("end position %v is out of bounds", edit.End)
 		}
-	// Splice the edit text in between the first and last lines of the edit.
-	prefix := string([]rune(content[edit.Start.Line])[:edit.Start.Column])
-	suffix := string([]rune(content[edit.End.Line])[edit.End.Column:])
-	newLines := strings.Split(prefix+edit.Text+suffix, "\n")
-	newContent := append(content[:edit.Start.Line], newLines...)
-	return append(newContent, content[edit.End.Line+1:]...), nil
+	}
+
+	var (
+		b            strings.Builder
+		line, column int
+	)
+	advance := func(toLine, toColumn int) {
+		for ; line < toLine; line++ {
+			b.WriteString(string([]rune(content[line])[column:]) + "\n")
+			column = 0
+		}
+		b.WriteString(string([]rune(content[line])[column:toColumn]))
+		column = toColumn
+	}
+	for _, edit := range newEdits {
+		advance(edit.Start.Line, edit.Start.Column)
+		b.WriteString(edit.Text)
+		line = edit.End.Line
+		column = edit.End.Column
+	}
+	advance(len(content)-1, len([]rune(content[len(content)-1])))
+	return strings.Split(b.String(), "\n"), nil
 }

internal/lsp/fake/edit_test.go

@@ -13,7 +13,7 @@ func TestApplyEdit(t *testing.T) {
 	tests := []struct {
 		label   string
 		content string
-		edit    Edit
+		edits   []Edit
 		want    string
 		wantErr bool
 	}{
@@ -23,57 +23,57 @@ func TestApplyEdit(t *testing.T) {
 		{
 			label:   "empty edit",
 			content: "hello",
-			edit:    Edit{},
+			edits:   []Edit{},
 			want:    "hello",
 		},
 		{
 			label:   "unicode edit",
 			content: "hello, 日本語",
-			edit: Edit{
+			edits: []Edit{{
 				Start: Pos{Line: 0, Column: 7},
 				End:   Pos{Line: 0, Column: 10},
 				Text:  "world",
-			},
+			}},
 			want: "hello, world",
 		},
 		{
 			label:   "range edit",
 			content: "ABC\nDEF\nGHI\nJKL",
-			edit: Edit{
+			edits: []Edit{{
 				Start: Pos{Line: 1, Column: 1},
 				End:   Pos{Line: 2, Column: 3},
 				Text:  "12\n345",
-			},
+			}},
 			want: "ABC\nD12\n345\nJKL",
 		},
 		{
 			label:   "end before start",
 			content: "ABC\nDEF\nGHI\nJKL",
-			edit: Edit{
+			edits: []Edit{{
 				End:   Pos{Line: 1, Column: 1},
 				Start: Pos{Line: 2, Column: 3},
 				Text:  "12\n345",
-			},
+			}},
 			wantErr: true,
 		},
 		{
 			label:   "out of bounds line",
 			content: "ABC\nDEF\nGHI\nJKL",
-			edit: Edit{
+			edits: []Edit{{
 				Start: Pos{Line: 1, Column: 1},
 				End:   Pos{Line: 4, Column: 3},
 				Text:  "12\n345",
-			},
+			}},
 			wantErr: true,
 		},
 		{
 			label:   "out of bounds column",
 			content: "ABC\nDEF\nGHI\nJKL",
-			edit: Edit{
+			edits: []Edit{{
 				Start: Pos{Line: 1, Column: 4},
 				End:   Pos{Line: 2, Column: 3},
 				Text:  "12\n345",
-			},
+			}},
 			wantErr: true,
 		},
 	}
@@ -82,7 +82,7 @@ func TestApplyEdit(t *testing.T) {
 		test := test
 		t.Run(test.label, func(t *testing.T) {
 			lines := strings.Split(test.content, "\n")
-			newLines, err := editContent(lines, test.edit)
+			newLines, err := editContent(lines, test.edits)
 			if (err != nil) != test.wantErr {
 				t.Errorf("got err %v, want error: %t", err, test.wantErr)
 			}

internal/lsp/fake/editor.go

@@ -235,6 +235,13 @@ func (e *Editor) CloseBuffer(ctx context.Context, path string) error {
 // SaveBuffer writes the content of the buffer specified by the given path to
 // the filesystem.
 func (e *Editor) SaveBuffer(ctx context.Context, path string) error {
+	if err := e.OrganizeImports(ctx, path); err != nil {
+		return fmt.Errorf("organizing imports before save: %v", err)
+	}
+	if err := e.FormatBuffer(ctx, path); err != nil {
+		return fmt.Errorf("formatting before save: %v", err)
+	}
+
 	e.mu.Lock()
 	buf, ok := e.buffers[path]
 	if !ok {
@@ -282,24 +289,22 @@ func (e *Editor) SaveBuffer(ctx context.Context, path string) error {
 
 // EditBuffer applies the given test edits to the buffer identified by path.
 func (e *Editor) EditBuffer(ctx context.Context, path string, edits []Edit) error {
-	params, err := e.doEdits(ctx, path, edits)
-	if err != nil {
-		return err
-	}
-	if e.server != nil {
-		if err := e.server.DidChange(ctx, params); err != nil {
-			return fmt.Errorf("DidChange: %v", err)
-		}
-	}
-	return nil
-}
-
-func (e *Editor) doEdits(ctx context.Context, path string, edits []Edit) (*protocol.DidChangeTextDocumentParams, error) {
 	e.mu.Lock()
 	defer e.mu.Unlock()
+	return e.editBufferLocked(ctx, path, edits)
+}
+
+// BufferText returns the content of the buffer with the given name.
+func (e *Editor) BufferText(name string) string {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	return e.buffers[name].text()
+}
+
+func (e *Editor) editBufferLocked(ctx context.Context, path string, edits []Edit) error {
 	buf, ok := e.buffers[path]
 	if !ok {
-		return nil, fmt.Errorf("unknown buffer %q", path)
+		return fmt.Errorf("unknown buffer %q", path)
 	}
 	var (
 		content = make([]string, len(buf.content))
@@ -307,16 +312,23 @@ func (e *Editor) doEdits(ctx context.Context, path string, edits []Edit) (*proto
 		evts    []protocol.TextDocumentContentChangeEvent
 	)
 	copy(content, buf.content)
-	for _, edit := range edits {
-		content, err = editContent(content, edit)
+	content, err = editContent(content, edits)
 	if err != nil {
-			return nil, err
-		}
-		evts = append(evts, edit.toProtocolChangeEvent())
+		return err
 	}
+
 	buf.content = content
 	buf.version++
 	e.buffers[path] = buf
+	// A simple heuristic: if there is only one edit, send it incrementally.
+	// Otherwise, send the entire content.
+	if len(edits) == 1 {
+		evts = append(evts, edits[0].toProtocolChangeEvent())
+	} else {
+		evts = append(evts, protocol.TextDocumentContentChangeEvent{
+			Text: buf.text(),
+		})
+	}
 	params := &protocol.DidChangeTextDocumentParams{
 		TextDocument: protocol.VersionedTextDocumentIdentifier{
 			Version: float64(buf.version),
@@ -326,7 +338,12 @@ func (e *Editor) doEdits(ctx context.Context, path string, edits []Edit) (*proto
 		},
 		ContentChanges: evts,
 	}
-	return params, nil
+	if e.server != nil {
+		if err := e.server.DidChange(ctx, params); err != nil {
+			return fmt.Errorf("DidChange: %v", err)
+		}
+	}
+	return nil
 }
 
 // GoToDefinition jumps to the definition of the symbol at the given position
@@ -354,6 +371,65 @@ func (e *Editor) GoToDefinition(ctx context.Context, path string, pos Pos) (stri
 	return newPath, newPos, nil
 }
 
+// OrganizeImports requests and performs the source.organizeImports codeAction.
+func (e *Editor) OrganizeImports(ctx context.Context, path string) error {
+	if e.server == nil {
+		return nil
+	}
+	params := &protocol.CodeActionParams{}
+	params.TextDocument.URI = e.ws.URI(path)
+
+	actions, err := e.server.CodeAction(ctx, params)
+	if err != nil {
+		return fmt.Errorf("textDocument/codeAction: %v", err)
+	}
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	for _, action := range actions {
+		if action.Kind == protocol.SourceOrganizeImports {
+			for _, change := range action.Edit.DocumentChanges {
+				path := e.ws.URIToPath(change.TextDocument.URI)
+				if float64(e.buffers[path].version) != change.TextDocument.Version {
+					// Skip edits for old versions.
+					continue
+				}
+				edits := convertEdits(change.Edits)
+				if err := e.editBufferLocked(ctx, path, edits); err != nil {
+					return fmt.Errorf("editing buffer %q: %v", path, err)
+				}
+			}
+		}
+	}
+	return nil
+}
+
+func convertEdits(protocolEdits []protocol.TextEdit) []Edit {
+	var edits []Edit
+	for _, lspEdit := range protocolEdits {
+		edits = append(edits, fromProtocolTextEdit(lspEdit))
+	}
+	return edits
+}
+
+// FormatBuffer gofmts a Go file.
+func (e *Editor) FormatBuffer(ctx context.Context, path string) error {
+	if e.server == nil {
+		return nil
+	}
+	// Because textDocument/formatting has no versions, we must block while
+	// performing formatting.
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	params := &protocol.DocumentFormattingParams{}
+	params.TextDocument.URI = e.ws.URI(path)
+	resp, err := e.server.Formatting(ctx, params)
+	if err != nil {
+		return fmt.Errorf("textDocument/formatting: %v", err)
+	}
+	edits := convertEdits(resp)
+	return e.editBufferLocked(ctx, path, edits)
+}
+
 func (e *Editor) checkBufferPosition(path string, pos Pos) error {
 	e.mu.Lock()
 	defer e.mu.Unlock()
@@ -366,6 +442,3 @@ func (e *Editor) checkBufferPosition(path string, pos Pos) error {
 	}
 	return nil
 }
-
-// TODO: expose more client functionality, for example Hover, CodeAction,
-// Rename, Completion, etc.  setting the content of an entire buffer, etc.

internal/lsp/fake/editor_test.go

@@ -23,20 +23,17 @@ func main() {
 `
 
 func TestClientEditing(t *testing.T) {
-	ws, err := NewWorkspace("test", []byte(exampleProgram))
+	ws, err := NewWorkspace("TestClientEditing", []byte(exampleProgram))
 	if err != nil {
 		t.Fatal(err)
 	}
 	defer ws.Close()
 	ctx := context.Background()
-	client := NewEditor(ws)
-	if err != nil {
+	editor := NewEditor(ws)
+	if err := editor.OpenFile(ctx, "main.go"); err != nil {
 		t.Fatal(err)
 	}
-	if err := client.OpenFile(ctx, "main.go"); err != nil {
-		t.Fatal(err)
-	}
-	if err := client.EditBuffer(ctx, "main.go", []Edit{
+	if err := editor.EditBuffer(ctx, "main.go", []Edit{
 		{
 			Start: Pos{5, 14},
 			End:   Pos{5, 26},
@@ -45,7 +42,7 @@ func TestClientEditing(t *testing.T) {
 	}); err != nil {
 		t.Fatal(err)
 	}
-	got := client.buffers["main.go"].text()
+	got := editor.buffers["main.go"].text()
 	want := `package main
 
 import "fmt"

internal/lsp/regtest/formatting_test.go

@@ -0,0 +1,93 @@
+package regtest
+
+import (
+	"context"
+	"testing"
+)
+
+const unformattedProgram = `
+-- main.go --
+package main
+import "fmt"
+func main(  ) {
+	fmt.Println("Hello World.")
+}
+-- main.go.golden --
+package main
+
+import "fmt"
+
+func main() {
+	fmt.Println("Hello World.")
+}
+`
+
+func TestFormatting(t *testing.T) {
+	runner.Run(t, unformattedProgram, func(ctx context.Context, t *testing.T, env *Env) {
+		env.OpenFile("main.go")
+		env.FormatBuffer("main.go")
+		got := env.E.BufferText("main.go")
+		want := env.ReadWorkspaceFile("main.go.golden")
+		if got != want {
+			t.Errorf("\n## got formatted file:\n%s\n## want:\n%s", got, want)
+		}
+	})
+}
+
+const disorganizedProgram = `
+-- main.go --
+package main
+
+import (
+	"fmt"
+	"errors"
+)
+func main(  ) {
+	fmt.Println(errors.New("bad"))
+}
+-- main.go.organized --
+package main
+
+import (
+	"errors"
+	"fmt"
+)
+func main(  ) {
+	fmt.Println(errors.New("bad"))
+}
+-- main.go.formatted --
+package main
+
+import (
+	"errors"
+	"fmt"
+)
+
+func main() {
+	fmt.Println(errors.New("bad"))
+}
+`
+
+func TestOrganizeImports(t *testing.T) {
+	runner.Run(t, disorganizedProgram, func(ctx context.Context, t *testing.T, env *Env) {
+		env.OpenFile("main.go")
+		env.OrganizeImports("main.go")
+		got := env.E.BufferText("main.go")
+		want := env.ReadWorkspaceFile("main.go.organized")
+		if got != want {
+			t.Errorf("\n## got formatted file:\n%s\n## want:\n%s", got, want)
+		}
+	})
+}
+
+func TestFormattingOnSave(t *testing.T) {
+	runner.Run(t, disorganizedProgram, func(ctx context.Context, t *testing.T, env *Env) {
+		env.OpenFile("main.go")
+		env.SaveBuffer("main.go")
+		got := env.E.BufferText("main.go")
+		want := env.ReadWorkspaceFile("main.go.formatted")
+		if got != want {
+			t.Errorf("\n## got formatted file:\n%s\n## want:\n%s", got, want)
+		}
+	})
+}

internal/lsp/regtest/wrappers.go

@@ -15,6 +15,17 @@ func (e *Env) RemoveFileFromWorkspace(name string) {
 	}
 }
 
+// ReadWorkspaceFile reads a file from the workspace, calling t.Fatal on any
+// error.
+func (e *Env) ReadWorkspaceFile(name string) string {
+	e.t.Helper()
+	content, err := e.W.ReadFile(name)
+	if err != nil {
+		e.t.Fatal(err)
+	}
+	return content
+}
+
 // OpenFile opens a file in the editor, calling t.Fatal on any error.
 func (e *Env) OpenFile(name string) {
 	e.t.Helper()
@@ -67,6 +78,23 @@ func (e *Env) GoToDefinition(name string, pos fake.Pos) (string, fake.Pos) {
 	return n, p
 }
 
+// FormatBuffer formats the editor buffer, calling t.Fatal on any error.
+func (e *Env) FormatBuffer(name string) {
+	e.t.Helper()
+	if err := e.E.FormatBuffer(e.ctx, name); err != nil {
+		e.t.Fatal(err)
+	}
+}
+
+// OrganizeImports processes the source.organizeImports codeAction, calling
+// t.Fatal on any error.
+func (e *Env) OrganizeImports(name string) {
+	e.t.Helper()
+	if err := e.E.OrganizeImports(e.ctx, name); err != nil {
+		e.t.Fatal(err)
+	}
+}
+
 // CloseEditor shuts down the editor, calling t.Fatal on any error.
 func (e *Env) CloseEditor() {
 	e.t.Helper()