gopls/internal/lsp: bundle certain quick-fixes with their diagnostic

To pragmatically avoid re-diagnosing the entire workspace, we can bundle
quick-fixes directly with their corresponding diagnostic, using the
Diagnostic.Data field added for this purpose in version 3.16 of the LSP
spec.

We should use this mechanism more generally, but for fixes with edits
we'd have to be careful that the edits are still valid in the current
snapshot. For now, be surgical.

This is the final regression we're tracking in the incremental gopls
issue (golang/go#57987).

Fixes golang/go#57987

Change-Id: Iaca91484e90341d677ecf573944edffef6e07255
Reviewed-on: https://go-review.googlesource.com/c/tools/+/497398
TryBot-Result: Gopher Robot <gobot@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>
Run-TryBot: Robert Findley <rfindley@google.com>
Reviewed-by: Alan Donovan <adonovan@google.com>

Author: findleyr

gopls/internal/lsp/cache/check.go

@@ -1542,14 +1542,18 @@ func depsErrors(ctx context.Context, m *source.Metadata, meta *metadataGraph, fs
 				if err != nil {
 					return nil, err
 				}
-				errors = append(errors, &source.Diagnostic{
+				diag := &source.Diagnostic{
 					URI:            imp.cgf.URI,
 					Range:          rng,
 					Severity:       protocol.SeverityError,
 					Source:         source.TypeError,
 					Message:        fmt.Sprintf("error while importing %v: %v", item, depErr.Err),
 					SuggestedFixes: fixes,
-				})
+				}
+				if !source.BundleQuickFixes(diag) {
+					bug.Reportf("failed to bundle fixes for diagnostic %q", diag.Message)
+				}
+				errors = append(errors, diag)
 			}
 		}
 	}
@@ -1585,14 +1589,18 @@ func depsErrors(ctx context.Context, m *source.Metadata, meta *metadataGraph, fs
 			if err != nil {
 				return nil, err
 			}
-			errors = append(errors, &source.Diagnostic{
+			diag := &source.Diagnostic{
 				URI:            pm.URI,
 				Range:          rng,
 				Severity:       protocol.SeverityError,
 				Source:         source.TypeError,
 				Message:        fmt.Sprintf("error while importing %v: %v", item, depErr.Err),
 				SuggestedFixes: fixes,
-			})
+			}
+			if !source.BundleQuickFixes(diag) {
+				bug.Reportf("failed to bundle fixes for diagnostic %q", diag.Message)
+			}
+			errors = append(errors, diag)
 			break
 		}
 	}

gopls/internal/lsp/code_action.go

@@ -473,6 +473,17 @@ func documentChanges(fh source.FileHandle, edits []protocol.TextEdit) []protocol
 
 func codeActionsMatchingDiagnostics(ctx context.Context, snapshot source.Snapshot, pdiags []protocol.Diagnostic, sdiags []*source.Diagnostic) ([]protocol.CodeAction, error) {
 	var actions []protocol.CodeAction
+	var unbundled []protocol.Diagnostic // diagnostics without bundled code actions in their Data field
+	for _, pd := range pdiags {
+		bundled := source.BundledQuickFixes(pd)
+		if len(bundled) > 0 {
+			actions = append(actions, bundled...)
+		} else {
+			// No bundled actions: keep searching for a match.
+			unbundled = append(unbundled, pd)
+		}
+	}
+
 	for _, sd := range sdiags {
 		var diag *protocol.Diagnostic
 		for _, pd := range pdiags {

gopls/internal/lsp/diagnostics.go

@@ -145,6 +145,9 @@ func computeDiagnosticHash(diags ...*source.Diagnostic) string {
 		fmt.Fprintf(h, "range: %s\n", d.Range)
 		fmt.Fprintf(h, "severity: %s\n", d.Severity)
 		fmt.Fprintf(h, "source: %s\n", d.Source)
+		if d.BundledFixes != nil {
+			fmt.Fprintf(h, "fixes: %s\n", *d.BundledFixes)
+		}
 	}
 	return fmt.Sprintf("%x", h.Sum(nil))
 }
@@ -771,6 +774,7 @@ func toProtocolDiagnostics(diagnostics []*source.Diagnostic) []protocol.Diagnost
 			Source:             string(diag.Source),
 			Tags:               emptySliceDiagnosticTag(diag.Tags),
 			RelatedInformation: diag.Related,
+			Data:               diag.BundledFixes,
 		}
 		if diag.Code != "" {
 			pdiag.Code = diag.Code

gopls/internal/lsp/protocol/generate/tables.go

@@ -68,6 +68,7 @@ var renameProp = map[prop]string{
 	{"Command", "arguments"}:       "[]json.RawMessage",
 	{"CompletionItem", "textEdit"}: "TextEdit",
 	{"Diagnostic", "code"}:         "interface{}",
+	{"Diagnostic", "data"}:         "json.RawMessage", // delay unmarshalling quickfixes
 
 	{"DocumentDiagnosticReportPartialResult", "relatedDocuments"}: "map[DocumentURI]interface{}",
 

gopls/internal/lsp/protocol/tsprotocol.go

@@ -6,7 +6,9 @@ package source
 
 import (
 	"context"
+	"encoding/json"
 
+	"golang.org/x/tools/gopls/internal/bug"
 	"golang.org/x/tools/gopls/internal/lsp/protocol"
 	"golang.org/x/tools/gopls/internal/span"
 )
@@ -136,3 +138,81 @@ func CombineDiagnostics(tdiags []*Diagnostic, adiags []*Diagnostic, outT, outA *
 
 	*outT = append(*outT, tdiags...)
 }
+
+// quickFixesJSON is a JSON-serializable list of quick fixes
+// to be saved in the protocol.Diagnostic.Data field.
+type quickFixesJSON struct {
+	// TODO(rfindley): pack some sort of identifier here for later
+	// lookup/validation?
+	Fixes []protocol.CodeAction
+}
+
+// BundleQuickFixes attempts to bundle sd.SuggestedFixes into the
+// sd.BundledFixes field, so that it can be round-tripped through the client.
+// It returns false if the quick-fixes cannot be bundled.
+func BundleQuickFixes(sd *Diagnostic) bool {
+	if len(sd.SuggestedFixes) == 0 {
+		return true
+	}
+	var actions []protocol.CodeAction
+	for _, fix := range sd.SuggestedFixes {
+		if fix.Edits != nil {
+			// For now, we only support bundled code actions that execute commands.
+			//
+			// In order to cleanly support bundled edits, we'd have to guarantee that
+			// the edits were generated on the current snapshot. But this naively
+			// implies that every fix would have to include a snapshot ID, which
+			// would require us to republish all diagnostics on each new snapshot.
+			//
+			// TODO(rfindley): in order to avoid this additional chatter, we'd need
+			// to build some sort of registry or other mechanism on the snapshot to
+			// check whether a diagnostic is still valid.
+			return false
+		}
+		action := protocol.CodeAction{
+			Title:   fix.Title,
+			Kind:    fix.ActionKind,
+			Command: fix.Command,
+		}
+		actions = append(actions, action)
+	}
+	fixes := quickFixesJSON{
+		Fixes: actions,
+	}
+	data, err := json.Marshal(fixes)
+	if err != nil {
+		bug.Reportf("marshalling quick fixes: %v", err)
+		return false
+	}
+	msg := json.RawMessage(data)
+	sd.BundledFixes = &msg
+	return true
+}
+
+// BundledQuickFixes extracts any bundled codeActions from the
+// diag.Data field.
+func BundledQuickFixes(diag protocol.Diagnostic) []protocol.CodeAction {
+	if diag.Data == nil {
+		return nil
+	}
+	var fix quickFixesJSON
+	if err := json.Unmarshal(*diag.Data, &fix); err != nil {
+		bug.Reportf("unmarshalling quick fix: %v", err)
+		return nil
+	}
+
+	var actions []protocol.CodeAction
+	for _, action := range fix.Fixes {
+		// See BundleQuickFixes: for now we only support bundling commands.
+		if action.Edit != nil {
+			bug.Reportf("bundled fix %q includes workspace edits", action.Title)
+			continue
+		}
+		// associate the action with the incoming diagnostic
+		// (Note that this does not mutate the fix.Fixes slice).
+		action.Diagnostics = []protocol.Diagnostic{diag}
+		actions = append(actions, action)
+	}
+
+	return actions
+}

gopls/internal/lsp/source/diagnostics.go

@@ -8,6 +8,7 @@ import (
 	"bytes"
 	"context"
 	"crypto/sha256"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"go/ast"
@@ -971,7 +972,18 @@ type Diagnostic struct {
 	Related []protocol.DiagnosticRelatedInformation
 
 	// Fields below are used internally to generate quick fixes. They aren't
-	// part of the LSP spec and don't leave the server.
+	// part of the LSP spec and historically didn't leave the server.
+	//
+	// Update(2023-05): version 3.16 of the LSP spec included support for the
+	// Diagnostic.data field, which holds arbitrary data preserved in the
+	// diagnostic for codeAction requests. This field allows bundling additional
+	// information for quick-fixes, and gopls can (and should) use this
+	// information to avoid re-evaluating diagnostics in code-action handlers.
+	//
+	// In order to stage this transition incrementally, the 'BundledFixes' field
+	// may store a 'bundled' (=json-serialized) form of the associated
+	// SuggestedFixes. Not all diagnostics have their fixes bundled.
+	BundledFixes   *json.RawMessage
 	SuggestedFixes []SuggestedFix
 }
 

gopls/internal/lsp/source/view.go

@@ -498,14 +498,8 @@ var _ = blah.Name
 			ReadDiagnostics("a/go.mod", &modDiags),
 		)
 
-		// golang.go#57987: now that gopls is incremental, we must be careful where
-		// we request diagnostics. We must design a simpler way to correlate
-		// published diagnostics with subsequent code action requests (see also the
-		// comment in Server.codeAction).
-		const canRequestCodeActionsForWorkspaceDiagnostics = false
-		if canRequestCodeActionsForWorkspaceDiagnostics {
-			env.ApplyQuickFixes("a/go.mod", modDiags.Diagnostics)
-			const want = `module mod.com
+		env.ApplyQuickFixes("a/go.mod", modDiags.Diagnostics)
+		const want = `module mod.com
 
 go 1.12
 
@@ -514,11 +508,10 @@ require (
 	example.com/blah/v2 v2.0.0
 )
 `
-			env.SaveBuffer("a/go.mod")
-			env.AfterChange(NoDiagnostics(ForFile("a/main.go")))
-			if got := env.BufferText("a/go.mod"); got != want {
-				t.Fatalf("suggested fixes failed:\n%s", compare.Text(want, got))
-			}
+		env.SaveBuffer("a/go.mod")
+		env.AfterChange(NoDiagnostics(ForFile("a/main.go")))
+		if got := env.BufferText("a/go.mod"); got != want {
+			t.Fatalf("suggested fixes failed:\n%s", compare.Text(want, got))
 		}
 	})
 }