internal/lsp: use versioned URIs in rename and code actions

This change adds support for returning versions along with file URIs, so
that the client can know when to apply changes. The version is not yet
propagated along to the internal/lsp/cache package, so this change will
have no effect (VS Code ignores a version of 0 and still applies the
changes).

A few minor changes made in the rename code (to remove the view
parameter). Some minor staticcheck fixes.

Updates golang/go#35243

Change-Id: Icc26bd9d9e5703c699f555424b94034c97b01d6f
Reviewed-on: https://go-review.googlesource.com/c/tools/+/206882
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
Reviewed-by: Ian Cottrell <iancottrell@google.com>

internal/lsp/cache/external.go

@@ -28,16 +28,16 @@ type nativeFileHandle struct {
 }
 
 func (fs *nativeFileSystem) GetFile(uri span.URI, kind source.FileKind) source.FileHandle {
-	version := "DOES NOT EXIST"
+	identifier := "DOES NOT EXIST"
 	if fi, err := os.Stat(uri.Filename()); err == nil {
-		version = fi.ModTime().String()
+		identifier = fi.ModTime().String()
 	}
 	return &nativeFileHandle{
 		fs: fs,
 		identity: source.FileIdentity{
-			URI:     uri,
-			Version: version,
-			Kind:    kind,
+			URI:        uri,
+			Identifier: identifier,
+			Kind:       kind,
 		},
 	}
 }

internal/lsp/cache/session.go

@@ -46,6 +46,7 @@ type overlay struct {
 	uri     span.URI
 	data    []byte
 	hash    string
+	version float64
 	kind    source.FileKind
 
 	// sameContentOnDisk is true if a file has been saved on disk,
@@ -424,9 +425,10 @@ func (o *overlay) FileSystem() source.FileSystem {
 
 func (o *overlay) Identity() source.FileIdentity {
 	return source.FileIdentity{
-		URI:     o.uri,
-		Version: o.hash,
-		Kind:    o.kind,
+		URI:        o.uri,
+		Identifier: o.hash,
+		Version:    o.version,
+		Kind:       o.kind,
 	}
 }
 func (o *overlay) Read(ctx context.Context) ([]byte, string, error) {

internal/lsp/cache/view.go

@@ -274,7 +274,7 @@ func (v *view) storeModFileVersions() {
 func (v *view) fileVersion(filename string, kind source.FileKind) string {
 	uri := span.FileURI(filename)
 	f := v.session.GetFile(uri, kind)
-	return f.Identity().Version
+	return f.Identity().Identifier
 }
 
 func (v *view) Shutdown(ctx context.Context) {

internal/lsp/cmd/imports.go

@@ -70,15 +70,19 @@ func (t *imports) Run(ctx context.Context, args ...string) error {
 	}
 	var edits []protocol.TextEdit
 	for _, a := range actions {
-		if a.Title == "Organize Imports" {
-			edits = (*a.Edit.Changes)[string(uri)]
+		if a.Title != "Organize Imports" {
+			continue
+		}
+		for _, c := range a.Edit.DocumentChanges {
+			if c.TextDocument.URI == string(uri) {
+				edits = append(edits, c.Edits...)
+			}
 		}
 	}
 	sedits, err := source.FromProtocolEdits(file.mapper, edits)
 	if err != nil {
 		return errors.Errorf("%v: %v", edits, err)
 	}
-
 	newContent := diff.ApplyEdits(string(file.mapper.Content), sedits)
 
 	filename := file.uri.Filename()

internal/lsp/cmd/rename.go

@@ -64,42 +64,39 @@ func (r *rename) Run(ctx context.Context, args ...string) error {
 	if file.err != nil {
 		return file.err
 	}
-
 	loc, err := file.mapper.Location(from)
 	if err != nil {
 		return err
 	}
-
 	p := protocol.RenameParams{
 		TextDocument: protocol.TextDocumentIdentifier{URI: loc.URI},
 		Position:     loc.Range.Start,
 		NewName:      args[1],
 	}
-	we, err := conn.Rename(ctx, &p)
+	edit, err := conn.Rename(ctx, &p)
 	if err != nil {
 		return err
 	}
-
-	// Make output order predictable
-	var keys []string
-	for u := range *we.Changes {
-		keys = append(keys, u)
+	var orderedURIs []string
+	edits := map[span.URI][]protocol.TextEdit{}
+	for _, c := range edit.DocumentChanges {
+		uri := span.URI(c.TextDocument.URI)
+		edits[uri] = append(edits[uri], c.Edits...)
+		orderedURIs = append(orderedURIs, c.TextDocument.URI)
 	}
-	sort.Strings(keys)
-	changeCount := len(keys)
+	sort.Strings(orderedURIs)
+	changeCount := len(orderedURIs)
 
-	for _, u := range keys {
-		edits := (*we.Changes)[u]
-		uri := span.NewURI(u)
+	for _, u := range orderedURIs {
+		uri := span.URI(u)
 		cmdFile := conn.AddFile(ctx, uri)
 		filename := cmdFile.uri.Filename()
 
 		// convert LSP-style edits to []diff.TextEdit cuz Spans are handy
-		renameEdits, err := source.FromProtocolEdits(cmdFile.mapper, edits)
+		renameEdits, err := source.FromProtocolEdits(cmdFile.mapper, edits[uri])
 		if err != nil {
 			return errors.Errorf("%v: %v", edits, err)
 		}
-
 		newContent := diff.ApplyEdits(string(cmdFile.mapper.Content), renameEdits)
 
 		switch {
@@ -114,7 +111,7 @@ func (r *rename) Run(ctx context.Context, args ...string) error {
 			diffs := diff.ToUnified(filename+".orig", filename, string(cmdFile.mapper.Content), renameEdits)
 			fmt.Print(diffs)
 		default:
-			if len(keys) > 1 {
+			if len(orderedURIs) > 1 {
 				fmt.Printf("%s:\n", filepath.Base(filename))
 			}
 			fmt.Print(string(newContent))

internal/lsp/cmd/suggested_fix.go

@@ -88,8 +88,13 @@ func (s *suggestedfix) Run(ctx context.Context, args ...string) error {
 	}
 	var edits []protocol.TextEdit
 	for _, a := range actions {
-		if a.IsPreferred || s.All {
-			edits = (*a.Edit.Changes)[string(uri)]
+		if !a.IsPreferred && !s.All {
+			continue
+		}
+		for _, c := range a.Edit.DocumentChanges {
+			if c.TextDocument.URI == string(uri) {
+				edits = append(edits, c.Edits...)
+			}
 		}
 	}
 

internal/lsp/code_action.go

@@ -28,12 +28,12 @@ func (s *Server) codeAction(ctx context.Context, params *protocol.CodeActionPara
 	}
 
 	snapshot := view.Snapshot()
+	fh := snapshot.Handle(ctx, f)
 
 	// Determine the supported actions for this file kind.
-	fileKind := snapshot.Handle(ctx, f).Identity().Kind
-	supportedCodeActions, ok := view.Options().SupportedCodeActions[fileKind]
+	supportedCodeActions, ok := view.Options().SupportedCodeActions[fh.Identity().Kind]
 	if !ok {
-		return nil, fmt.Errorf("no supported code actions for %v file kind", fileKind)
+		return nil, fmt.Errorf("no supported code actions for %v file kind", fh.Identity().Kind)
 	}
 
 	// The Only field of the context specifies which code actions the client wants.
@@ -52,7 +52,7 @@ func (s *Server) codeAction(ctx context.Context, params *protocol.CodeActionPara
 	}
 
 	var codeActions []protocol.CodeAction
-	switch fileKind {
+	switch fh.Identity().Kind {
 	case source.Mod:
 		if !wanted[protocol.SourceOrganizeImports] {
 			return nil, nil
@@ -92,9 +92,7 @@ func (s *Server) codeAction(ctx context.Context, params *protocol.CodeActionPara
 							Title: importFixTitle(importFix.Fix),
 							Kind:  protocol.QuickFix,
 							Edit: &protocol.WorkspaceEdit{
-								Changes: &map[string][]protocol.TextEdit{
-									string(uri): importFix.Edits,
-								},
+								DocumentChanges: documentChanges(fh, importFix.Edits),
 							},
 							Diagnostics: fixDiagnostics,
 						})
@@ -107,9 +105,7 @@ func (s *Server) codeAction(ctx context.Context, params *protocol.CodeActionPara
 				Title: "Organize Imports",
 				Kind:  protocol.SourceOrganizeImports,
 				Edit: &protocol.WorkspaceEdit{
-					Changes: &map[string][]protocol.TextEdit{
-						string(uri): edits,
-					},
+					DocumentChanges: documentChanges(fh, edits),
 				},
 			})
 		}
@@ -223,19 +219,37 @@ func quickFixes(ctx context.Context, s source.Snapshot, f source.File, diagnosti
 			continue
 		}
 		for _, fix := range srcErr.SuggestedFixes {
-			edits := make(map[string][]protocol.TextEdit)
-			for uri, e := range fix.Edits {
-				edits[protocol.NewURI(uri)] = e
-			}
-			codeActions = append(codeActions, protocol.CodeAction{
+			action := protocol.CodeAction{
 				Title:       fix.Title,
 				Kind:        protocol.QuickFix,
 				Diagnostics: []protocol.Diagnostic{diag},
-				Edit: &protocol.WorkspaceEdit{
-					Changes: &edits,
-				},
-			})
+				Edit:        &protocol.WorkspaceEdit{},
+			}
+			for uri, edits := range fix.Edits {
+				f, err := s.View().GetFile(ctx, uri)
+				if err != nil {
+					log.Error(ctx, "no file", err, telemetry.URI.Of(uri))
+					continue
+				}
+				fh := s.Handle(ctx, f)
+				action.Edit.DocumentChanges = append(action.Edit.DocumentChanges, documentChanges(fh, edits)...)
+			}
+			codeActions = append(codeActions, action)
 		}
 	}
 	return codeActions, nil
 }
+
+func documentChanges(fh source.FileHandle, edits []protocol.TextEdit) []protocol.TextDocumentEdit {
+	return []protocol.TextDocumentEdit{
+		{
+			TextDocument: protocol.VersionedTextDocumentIdentifier{
+				Version: fh.Identity().Version,
+				TextDocumentIdentifier: protocol.TextDocumentIdentifier{
+					URI: protocol.NewURI(fh.Identity().URI),
+				},
+			},
+			Edits: edits,
+		},
+	}
+}

internal/lsp/lsp_test.go

@@ -310,17 +310,14 @@ func (r *runner) Import(t *testing.T, spn span.Span) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	var edits []protocol.TextEdit
-	for _, a := range actions {
-		if a.Title == "Organize Imports" {
-			edits = (*a.Edit.Changes)[string(uri)]
+	got := string(m.Content)
+	if len(actions) > 0 {
+		res, err := applyWorkspaceEdits(r, actions[0].Edit)
+		if err != nil {
+			t.Fatal(err)
 		}
+		got = res[uri]
 	}
-	sedits, err := source.FromProtocolEdits(m, edits)
-	if err != nil {
-		t.Error(err)
-	}
-	got := diff.ApplyEdits(string(m.Content), sedits)
 	if goimported != got {
 		t.Errorf("import failed for %s, expected:\n%v\ngot:\n%v", filename, goimported, got)
 	}
@@ -347,25 +344,18 @@ func (r *runner) SuggestedFix(t *testing.T, spn span.Span) {
 			Diagnostics: toProtocolDiagnostics(r.ctx, diagnostics[uri]),
 		},
 	})
-	if err != nil {
-		t.Error(err)
-		return
-	}
-	m, err := r.data.Mapper(f.URI())
 	if err != nil {
 		t.Fatal(err)
 	}
-	var edits []protocol.TextEdit
-	for _, a := range actions {
-		if a.Title == "Remove" {
-			edits = (*a.Edit.Changes)[string(uri)]
-		}
+	// TODO: This test should probably be able to handle multiple code actions.
+	if len(actions) > 1 {
+		t.Fatal("expected only 1 code action")
 	}
-	sedits, err := source.FromProtocolEdits(m, edits)
+	res, err := applyWorkspaceEdits(r, actions[0].Edit)
 	if err != nil {
-		t.Error(err)
+		t.Fatal(err)
 	}
-	got := diff.ApplyEdits(string(m.Content), sedits)
+	got := res[uri]
 	fixed := string(r.data.Golden("suggestedfix", filename, func() ([]byte, error) {
 		return []byte(got), nil
 	}))
@@ -563,7 +553,7 @@ func (r *runner) Rename(t *testing.T, spn span.Span, newText string) {
 		t.Fatalf("failed for %v: %v", spn, err)
 	}
 
-	workspaceEdits, err := r.server.Rename(r.ctx, &protocol.RenameParams{
+	wedit, err := r.server.Rename(r.ctx, &protocol.RenameParams{
 		TextDocument: protocol.TextDocumentIdentifier{
 			URI: protocol.NewURI(uri),
 		},
@@ -579,33 +569,26 @@ func (r *runner) Rename(t *testing.T, spn span.Span, newText string) {
 		}
 		return
 	}
-
-	var res []string
-	for uri, edits := range *workspaceEdits.Changes {
-		m, err := r.data.Mapper(span.URI(uri))
-		if err != nil {
-			t.Fatal(err)
-		}
-		sedits, err := source.FromProtocolEdits(m, edits)
-		if err != nil {
-			t.Error(err)
-		}
-		filename := filepath.Base(m.URI.Filename())
-		contents := applyEdits(string(m.Content), sedits)
-		if len(*workspaceEdits.Changes) > 1 {
-			contents = fmt.Sprintf("%s:\n%s", filename, contents)
-		}
-		res = append(res, contents)
+	res, err := applyWorkspaceEdits(r, wedit)
+	if err != nil {
+		t.Fatal(err)
 	}
-
-	// Sort on filename
-	sort.Strings(res)
+	var orderedURIs []string
+	for uri := range res {
+		orderedURIs = append(orderedURIs, string(uri))
+	}
+	sort.Strings(orderedURIs)
 
 	var got string
-	for i, val := range res {
+	for i := 0; i < len(res); i++ {
 		if i != 0 {
 			got += "\n"
 		}
+		uri := span.URI(orderedURIs[i])
+		if len(res) > 1 {
+			got += filepath.Base(uri.Filename()) + ":\n"
+		}
+		val := res[uri]
 		got += val
 	}
 
@@ -650,6 +633,24 @@ func (r *runner) PrepareRename(t *testing.T, src span.Span, want *source.Prepare
 	}
 }
 
+func applyWorkspaceEdits(r *runner, wedit *protocol.WorkspaceEdit) (map[span.URI]string, error) {
+	res := map[span.URI]string{}
+	for _, docEdits := range wedit.DocumentChanges {
+		uri := span.URI(docEdits.TextDocument.URI)
+		m, err := r.data.Mapper(uri)
+		if err != nil {
+			return nil, err
+		}
+		res[uri] = string(m.Content)
+		sedits, err := source.FromProtocolEdits(m, docEdits.Edits)
+		if err != nil {
+			return nil, err
+		}
+		res[uri] = applyEdits(res[uri], sedits)
+	}
+	return res, nil
+}
+
 func applyEdits(contents string, edits []diff.TextEdit) string {
 	res := contents