internal/lsp: organize imports as a quickfix

This change adds the ability to organize all imports as quickfix on a
diagnostic, in addition to the standard "source.organizeImports" code
action kind. This change still requires a bit of work, since really, we
should only add one import per diagnostic and each edit should be the
edits necessary for the addition or removal of the individual import,
but this is a good start.

Change-Id: I57679446bec833d215c35533240102fb61c86f20
Reviewed-on: https://go-review.googlesource.com/c/tools/+/172397
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 177c145..c49b19e 100644
--- a/internal/lsp/code_action.go
+++ b/internal/lsp/code_action.go
@@ -28,12 +28,29 @@
 	var codeActions []protocol.CodeAction
 	// Determine what code actions we should take based on the diagnostics.
 	if findImportErrors(params.Context.Diagnostics) {
-		codeAction, err := organizeImports(ctx, view, spn)
+		edits, err := organizeImports(ctx, view, spn)
 		if err != nil {
 			return nil, err
 		}
-		if codeAction != nil {
-			codeActions = append(codeActions, *codeAction)
+		if len(edits) > 0 {
+			// 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.
+			codeActions = append(codeActions, protocol.CodeAction{
+				Title: "Organize Imports",
+				Kind:  protocol.SourceOrganizeImports,
+				Edit: &protocol.WorkspaceEdit{
+					Changes: &map[string][]protocol.TextEdit{
+						string(spn.URI()): edits,
+					},
+				},
+			})
+			// Add any quick fixes for each import-related diagnostic that we see.
+			fixes, err := quickFixes(spn.URI(), params.Context.Diagnostics, edits)
+			if err != nil {
+				return nil, err
+			}
+			codeActions = append(codeActions, fixes...)
 		}
 	}
 	return codeActions, nil
@@ -60,7 +77,7 @@
 	return false
 }
 
-func organizeImports(ctx context.Context, v source.View, s span.Span) (*protocol.CodeAction, error) {
+func organizeImports(ctx context.Context, v source.View, s span.Span) ([]protocol.TextEdit, error) {
 	f, m, err := newColumnMap(ctx, v, s.URI())
 	if err != nil {
 		return nil, err
@@ -81,21 +98,22 @@
 	if err != nil {
 		return nil, err
 	}
-	protocolEdits, err := ToProtocolEdits(m, edits)
-	if err != nil {
-		return nil, err
-	}
-	if len(protocolEdits) == 0 {
-		return nil, nil
-	}
-	codeAction := protocol.CodeAction{
-		Title: "Organize Imports",
-		Kind:  protocol.SourceOrganizeImports,
-		Edit: &protocol.WorkspaceEdit{
-			Changes: &map[string][]protocol.TextEdit{
-				string(s.URI()): protocolEdits,
+	return ToProtocolEdits(m, edits)
+}
+
+// TODO(rstambler): Separate this into a set of codeActions per diagnostic,
+// where each action is the addition or removal of one import.
+// This can only be done when https://golang.org/issue/31493 is resolved.
+func quickFixes(uri span.URI, diagnostics []protocol.Diagnostic, edits []protocol.TextEdit) ([]protocol.CodeAction, error) {
+	return []protocol.CodeAction{
+		{
+			Title: "Organize All Imports",
+			Kind:  protocol.QuickFix,
+			Edit: &protocol.WorkspaceEdit{
+				Changes: &map[string][]protocol.TextEdit{
+					string(uri): edits,
+				},
 			},
 		},
-	}
-	return &codeAction, nil
+	}, nil
 }