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>
diff --git a/internal/lsp/cache/external.go b/internal/lsp/cache/external.go
index a75ec07..2498bf6 100644
--- a/internal/lsp/cache/external.go
+++ b/internal/lsp/cache/external.go
@@ -28,16 +28,16 @@
 }
 
 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,
 		},
 	}
 }
diff --git a/internal/lsp/cache/session.go b/internal/lsp/cache/session.go
index 2174df4..e81ee94 100644
--- a/internal/lsp/cache/session.go
+++ b/internal/lsp/cache/session.go
@@ -46,6 +46,7 @@
 	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) 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) {
diff --git a/internal/lsp/cache/view.go b/internal/lsp/cache/view.go
index 5bd8d10..03a10ac 100644
--- a/internal/lsp/cache/view.go
+++ b/internal/lsp/cache/view.go
@@ -274,7 +274,7 @@
 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) {
diff --git a/internal/lsp/cmd/imports.go b/internal/lsp/cmd/imports.go
index 19531df..2127e25 100644
--- a/internal/lsp/cmd/imports.go
+++ b/internal/lsp/cmd/imports.go
@@ -70,15 +70,19 @@
 	}
 	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()
diff --git a/internal/lsp/cmd/rename.go b/internal/lsp/cmd/rename.go
index 47e26be..98f7d69 100644
--- a/internal/lsp/cmd/rename.go
+++ b/internal/lsp/cmd/rename.go
@@ -64,42 +64,39 @@
 	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 @@
 			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))
diff --git a/internal/lsp/cmd/suggested_fix.go b/internal/lsp/cmd/suggested_fix.go
index 25ee8b6..15c6bc9 100644
--- a/internal/lsp/cmd/suggested_fix.go
+++ b/internal/lsp/cmd/suggested_fix.go
@@ -88,8 +88,13 @@
 	}
 	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...)
+			}
 		}
 	}
 
diff --git a/internal/lsp/code_action.go b/internal/lsp/code_action.go
index 0d39661..cf66505 100644
--- a/internal/lsp/code_action.go
+++ b/internal/lsp/code_action.go
@@ -28,12 +28,12 @@
 	}
 
 	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 @@
 	}
 
 	var codeActions []protocol.CodeAction
-	switch fileKind {
+	switch fh.Identity().Kind {
 	case source.Mod:
 		if !wanted[protocol.SourceOrganizeImports] {
 			return nil, nil
@@ -92,9 +92,7 @@
 							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 @@
 				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 @@
 			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,
+		},
+	}
+}
diff --git a/internal/lsp/lsp_test.go b/internal/lsp/lsp_test.go
index 8e2615b..f20b967 100644
--- a/internal/lsp/lsp_test.go
+++ b/internal/lsp/lsp_test.go
@@ -310,17 +310,14 @@
 	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)
 	}
@@ -348,24 +345,17 @@
 		},
 	})
 	if err != nil {
-		t.Error(err)
-		return
+		t.Fatal(err)
 	}
-	m, err := r.data.Mapper(f.URI())
+	// TODO: This test should probably be able to handle multiple code actions.
+	if len(actions) > 1 {
+		t.Fatal("expected only 1 code action")
+	}
+	res, err := applyWorkspaceEdits(r, actions[0].Edit)
 	if err != nil {
 		t.Fatal(err)
 	}
-	var edits []protocol.TextEdit
-	for _, a := range actions {
-		if a.Title == "Remove" {
-			edits = (*a.Edit.Changes)[string(uri)]
-		}
-	}
-	sedits, err := source.FromProtocolEdits(m, edits)
-	if err != nil {
-		t.Error(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 @@
 		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 @@
 		}
 		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 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
 
diff --git a/internal/lsp/rename.go b/internal/lsp/rename.go
index 9ce3672..9664b5c 100644
--- a/internal/lsp/rename.go
+++ b/internal/lsp/rename.go
@@ -23,16 +23,22 @@
 	if err != nil {
 		return nil, err
 	}
-	edits, err := ident.Rename(ctx, view, params.NewName)
+	edits, err := ident.Rename(ctx, params.NewName)
 	if err != nil {
 		return nil, err
 	}
-	changes := make(map[string][]protocol.TextEdit)
+	var docChanges []protocol.TextDocumentEdit
 	for uri, e := range edits {
-		changes[protocol.NewURI(uri)] = e
+		f, err := view.GetFile(ctx, uri)
+		if err != nil {
+			return nil, err
+		}
+		fh := ident.Snapshot.Handle(ctx, f)
+		docChanges = append(docChanges, documentChanges(fh, e)...)
 	}
-
-	return &protocol.WorkspaceEdit{Changes: &changes}, nil
+	return &protocol.WorkspaceEdit{
+		DocumentChanges: docChanges,
+	}, nil
 }
 
 func (s *Server) prepareRename(ctx context.Context, params *protocol.PrepareRenameParams) (*protocol.Range, error) {
@@ -42,11 +48,15 @@
 	if err != nil {
 		return nil, err
 	}
+	ident, err := source.Identifier(ctx, view, f, params.Position)
+	if err != nil {
+		return nil, nil // ignore errors
+	}
 	// Do not return errors here, as it adds clutter.
 	// Returning a nil result means there is not a valid rename.
-	item, err := source.PrepareRename(ctx, view, f, params.Position)
+	item, err := ident.PrepareRename(ctx)
 	if err != nil {
-		return nil, nil
+		return nil, nil // ignore errors
 	}
 	// TODO(suzmue): return ident.Name as the placeholder text.
 	return &item.Range, nil
diff --git a/internal/lsp/source/diagnostics.go b/internal/lsp/source/diagnostics.go
index fe4d38a..e85c623 100644
--- a/internal/lsp/source/diagnostics.go
+++ b/internal/lsp/source/diagnostics.go
@@ -109,6 +109,7 @@
 
 func diagnostics(ctx context.Context, view View, pkg Package, reports map[span.URI][]Diagnostic) bool {
 	ctx, done := trace.StartSpan(ctx, "source.diagnostics", telemetry.Package.Of(pkg.ID()))
+	_ = ctx // circumvent SA4006
 	defer done()
 
 	diagSets := make(map[span.URI]*diagnosticSet)
diff --git a/internal/lsp/source/implementation.go b/internal/lsp/source/implementation.go
index 2793333..f82f646 100644
--- a/internal/lsp/source/implementation.go
+++ b/internal/lsp/source/implementation.go
@@ -83,7 +83,7 @@
 			}
 		}
 		if containingFile == nil {
-			return nil, fmt.Errorf("Failed to find file %q in package %v", file.Name(), pkg.PkgPath())
+			return nil, fmt.Errorf("failed to find file %q in package %v", file.Name(), pkg.PkgPath())
 		}
 
 		uri := containingFile.Identity().URI
diff --git a/internal/lsp/source/rename.go b/internal/lsp/source/rename.go
index d6b50f8..1717c66 100644
--- a/internal/lsp/source/rename.go
+++ b/internal/lsp/source/rename.go
@@ -41,15 +41,10 @@
 	Text  string
 }
 
-func PrepareRename(ctx context.Context, view View, f File, pos protocol.Position) (*PrepareItem, error) {
+func (i *IdentifierInfo) PrepareRename(ctx context.Context) (*PrepareItem, error) {
 	ctx, done := trace.StartSpan(ctx, "source.PrepareRename")
 	defer done()
 
-	i, err := Identifier(ctx, view, f, pos)
-	if err != nil {
-		return nil, err
-	}
-
 	// TODO(rstambler): We should handle this in a better way.
 	// If the object declaration is nil, assume it is an import spec.
 	if i.Declaration.obj == nil {
@@ -77,7 +72,7 @@
 }
 
 // Rename returns a map of TextEdits for each file modified when renaming a given identifier within a package.
-func (i *IdentifierInfo) Rename(ctx context.Context, view View, newName string) (map[span.URI][]protocol.TextEdit, error) {
+func (i *IdentifierInfo) Rename(ctx context.Context, newName string) (map[span.URI][]protocol.TextEdit, error) {
 	ctx, done := trace.StartSpan(ctx, "source.Rename")
 	defer done()
 
@@ -90,7 +85,7 @@
 		if err != nil {
 			return nil, err
 		}
-		return ident.Rename(ctx, view, newName)
+		return ident.Rename(ctx, newName)
 	}
 	if i.Name == newName {
 		return nil, errors.Errorf("old and new names are the same: %s", newName)
@@ -117,7 +112,7 @@
 
 	r := renamer{
 		ctx:          ctx,
-		fset:         view.Session().Cache().FileSet(),
+		fset:         i.Snapshot.View().Session().Cache().FileSet(),
 		refs:         refs,
 		objsToUpdate: make(map[types.Object]bool),
 		from:         i.Name,
@@ -147,7 +142,7 @@
 	for uri, edits := range changes {
 		// These edits should really be associated with FileHandles for maximal correctness.
 		// For now, this is good enough.
-		f, err := view.GetFile(ctx, uri)
+		f, err := i.Snapshot.View().GetFile(ctx, uri)
 		if err != nil {
 			return nil, err
 		}
diff --git a/internal/lsp/source/source_test.go b/internal/lsp/source/source_test.go
index 8e550f6..4281635 100644
--- a/internal/lsp/source/source_test.go
+++ b/internal/lsp/source/source_test.go
@@ -663,7 +663,7 @@
 		t.Error(err)
 		return
 	}
-	changes, err := ident.Rename(r.ctx, r.view, newText)
+	changes, err := ident.Rename(r.ctx, newText)
 	if err != nil {
 		renamed := string(r.data.Golden(tag, spn.URI().Filename(), func() ([]byte, error) {
 			return []byte(err.Error()), nil
@@ -747,7 +747,14 @@
 		t.Fatal(err)
 	}
 	// Find the identifier at the position.
-	item, err := source.PrepareRename(ctx, r.view, f, srcRng.Start)
+	ident, err := source.Identifier(ctx, r.view, f, srcRng.Start)
+	if err != nil {
+		if want.Text != "" { // expected an ident.
+			t.Errorf("prepare rename failed for %v: got error: %v", src, err)
+		}
+		return
+	}
+	item, err := ident.PrepareRename(ctx)
 	if err != nil {
 		if want.Text != "" { // expected an ident.
 			t.Errorf("prepare rename failed for %v: got error: %v", src, err)
diff --git a/internal/lsp/source/view.go b/internal/lsp/source/view.go
index 842e5e2..568151a 100644
--- a/internal/lsp/source/view.go
+++ b/internal/lsp/source/view.go
@@ -20,13 +20,21 @@
 
 // FileIdentity uniquely identifies a file at a version from a FileSystem.
 type FileIdentity struct {
-	URI     span.URI
-	Version string
-	Kind    FileKind
+	URI span.URI
+
+	// Version is the version of the file, as specified by the client.
+	Version float64
+
+	// Identifier represents a unique identifier for the file.
+	// It could be a file's modification time or its SHA1 hash if it is not on disk.
+	Identifier string
+
+	// Kind is the file's kind.
+	Kind FileKind
 }
 
 func (identity FileIdentity) String() string {
-	return fmt.Sprintf("%s%s%s", identity.URI, identity.Version, identity.Kind)
+	return fmt.Sprintf("%s%f%s%s", identity.URI, identity.Version, identity.Identifier, identity.Kind)
 }
 
 // FileHandle represents a handle to a specific version of a single file from