internal/lsp: handle the context.only parameter for code actions

This change refactors code actions to handle the Context.Only parameter,
which indicates which code actions a language server should execute.

Change-Id: Iddfccbbeba3a53fde2aa8df844434f2ab9d01666
Reviewed-on: https://go-review.googlesource.com/c/tools/+/184158
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
Reviewed-by: Ian Cottrell <iancottrell@google.com>
diff --git a/internal/lsp/code_action.go b/internal/lsp/code_action.go
index ea4607a..e5195e7 100644
--- a/internal/lsp/code_action.go
+++ b/internal/lsp/code_action.go
@@ -6,6 +6,7 @@
 
 import (
 	"context"
+	"fmt"
 	"strings"
 
 	"golang.org/x/tools/internal/lsp/protocol"
@@ -14,7 +15,24 @@
 )
 
 func (s *Server) codeAction(ctx context.Context, params *protocol.CodeActionParams) ([]protocol.CodeAction, error) {
+
+	// The Only field of the context specifies which code actions the client wants.
+	// If Only is empty, assume that the client wants all of the possible code actions.
+	var wanted map[protocol.CodeActionKind]bool
+	if len(params.Context.Only) == 0 {
+		wanted = s.supportedCodeActions
+	} else {
+		wanted = make(map[protocol.CodeActionKind]bool)
+		for _, only := range params.Context.Only {
+			wanted[only] = s.supportedCodeActions[only]
+		}
+	}
+
 	uri := span.NewURI(params.TextDocument.URI)
+	if len(wanted) == 0 {
+		return nil, fmt.Errorf("no supported code action to execute for %s, wanted %v", uri, params.Context.Only)
+	}
+
 	view := s.session.ViewOf(uri)
 	gof, m, err := getGoFile(ctx, view, uri)
 	if err != nil {
@@ -24,25 +42,27 @@
 	if err != nil {
 		return nil, err
 	}
+
 	var codeActions []protocol.CodeAction
-	// TODO(rstambler): Handle params.Context.Only when VSCode-Go uses a
-	// version of vscode-languageclient that fixes
-	// https://github.com/Microsoft/vscode-languageserver-node/issues/442.
+
 	edits, err := organizeImports(ctx, view, spn)
 	if err != nil {
 		return nil, err
 	}
-	if len(edits) > 0 {
-		codeActions = append(codeActions, protocol.CodeAction{
-			Title: "Organize Imports",
-			Kind:  protocol.SourceOrganizeImports,
-			Edit: &protocol.WorkspaceEdit{
-				Changes: &map[string][]protocol.TextEdit{
-					string(spn.URI()): edits,
-				},
-			},
-		})
-		// If we also have diagnostics, we can associate them with quick fixes.
+
+	// If the user wants to see quickfixes.
+	if wanted[protocol.QuickFix] {
+		// First, add the quick fixes reported by go/analysis.
+		// TODO: Enable this when this actually works. For now, it's needless work.
+		if s.wantSuggestedFixes {
+			qf, err := quickFixes(ctx, view, gof)
+			if err != nil {
+				view.Session().Logger().Errorf(ctx, "quick fixes failed for %s: %v", uri, err)
+			}
+			codeActions = append(codeActions, qf...)
+		}
+
+		// If we also have diagnostics for missing imports, we can associate them with quick fixes.
 		if findImportErrors(params.Context.Diagnostics) {
 			// TODO(rstambler): Separate this into a set of codeActions per diagnostic,
 			// where each action is the addition or removal of one import.
@@ -57,26 +77,21 @@
 				},
 			})
 		}
-		diags := gof.GetPackage(ctx).GetDiagnostics()
-		for _, diag := range diags {
-			pdiag, err := toProtocolDiagnostic(ctx, view, diag)
-			if err != nil {
-				return nil, err
-			}
-			for _, ca := range diag.SuggestedFixes {
-				codeActions = append(codeActions, protocol.CodeAction{
-					Title: ca.Title,
-					Kind:  protocol.QuickFix, // TODO(matloob): Be more accurate about these?
-					Edit: &protocol.WorkspaceEdit{
-						Changes: &map[string][]protocol.TextEdit{
-							string(spn.URI()): edits,
-						},
-					},
-					Diagnostics: []protocol.Diagnostic{pdiag},
-				})
-			}
-		}
 	}
+
+	// Add the results of import organization as source.OrganizeImports.
+	if wanted[protocol.SourceOrganizeImports] {
+		codeActions = append(codeActions, protocol.CodeAction{
+			Title: "Organize Imports",
+			Kind:  protocol.SourceOrganizeImports,
+			Edit: &protocol.WorkspaceEdit{
+				Changes: &map[string][]protocol.TextEdit{
+					string(spn.URI()): edits,
+				},
+			},
+		})
+	}
+
 	return codeActions, nil
 }
 
@@ -112,3 +127,39 @@
 	}
 	return false
 }
+
+func quickFixes(ctx context.Context, view source.View, gof source.GoFile) ([]protocol.CodeAction, error) {
+	var codeActions []protocol.CodeAction
+
+	// TODO: This is technically racy because the diagnostics provided by the code action
+	// may not be the same as the ones that gopls is aware of.
+	// We need to figure out some way to solve this problem.
+	diags := gof.GetPackage(ctx).GetDiagnostics()
+	for _, diag := range diags {
+		pdiag, err := toProtocolDiagnostic(ctx, view, diag)
+		if err != nil {
+			return nil, err
+		}
+		for _, ca := range diag.SuggestedFixes {
+			_, m, err := getGoFile(ctx, view, diag.URI())
+			if err != nil {
+				return nil, err
+			}
+			edits, err := ToProtocolEdits(m, ca.Edits)
+			if err != nil {
+				return nil, err
+			}
+			codeActions = append(codeActions, protocol.CodeAction{
+				Title: ca.Title,
+				Kind:  protocol.QuickFix, // TODO(matloob): Be more accurate about these?
+				Edit: &protocol.WorkspaceEdit{
+					Changes: &map[string][]protocol.TextEdit{
+						string(diag.URI()): edits,
+					},
+				},
+				Diagnostics: []protocol.Diagnostic{pdiag},
+			})
+		}
+	}
+	return codeActions, nil
+}
diff --git a/internal/lsp/general.go b/internal/lsp/general.go
index c976760..2ba32b4 100644
--- a/internal/lsp/general.go
+++ b/internal/lsp/general.go
@@ -34,6 +34,11 @@
 		}
 	}
 
+	s.supportedCodeActions = map[protocol.CodeActionKind]bool{
+		protocol.SourceOrganizeImports: true,
+		protocol.QuickFix:              true,
+	}
+
 	s.setClientCapabilities(params.Capabilities)
 
 	folders := params.WorkspaceFolders
@@ -188,10 +193,14 @@
 	if usePlaceholders, ok := c["usePlaceholders"].(bool); ok {
 		s.usePlaceholders = usePlaceholders
 	}
-	// Check if user has disabled documentation on hover.
+	// Check if the user has disabled documentation on hover.
 	if noDocsOnHover, ok := c["noDocsOnHover"].(bool); ok {
 		s.noDocsOnHover = noDocsOnHover
 	}
+	// Check if the user wants to see suggested fixes from go/analysis.
+	if wantSuggestedFixes, ok := c["wantSuggestedFixes"].(bool); ok {
+		s.wantSuggestedFixes = wantSuggestedFixes
+	}
 	// Check if the user has explicitly disabled any analyses.
 	if disabledAnalyses, ok := c["experimentalDisabledAnalyses"].([]interface{}); ok {
 		s.disabledAnalyses = make(map[string]struct{})
diff --git a/internal/lsp/lsp_test.go b/internal/lsp/lsp_test.go
index 8f3eda4..39603a7 100644
--- a/internal/lsp/lsp_test.go
+++ b/internal/lsp/lsp_test.go
@@ -51,6 +51,10 @@
 		server: &Server{
 			session:     session,
 			undelivered: make(map[span.URI][]source.Diagnostic),
+			supportedCodeActions: map[protocol.CodeActionKind]bool{
+				protocol.SourceOrganizeImports: true,
+				protocol.QuickFix:              true,
+			},
 		},
 		data: data,
 	}
diff --git a/internal/lsp/server.go b/internal/lsp/server.go
index 97f8b5b..44c145b 100644
--- a/internal/lsp/server.go
+++ b/internal/lsp/server.go
@@ -78,6 +78,9 @@
 	dynamicConfigurationSupported bool
 	preferredContentFormat        protocol.MarkupKind
 	disabledAnalyses              map[string]struct{}
+	wantSuggestedFixes            bool
+
+	supportedCodeActions map[protocol.CodeActionKind]bool
 
 	textDocumentSyncKind protocol.TextDocumentSyncKind