From 57eff0d8ac5e8e0e378638d379d584e8f5029ea5 Mon Sep 17 00:00:00 2001
From: Rebecca Stambler <rstambler@golang.org>
Date: Fri, 14 Dec 2018 17:00:24 -0500
Subject: [PATCH] internal/lsp: add support for running goimports as a code
 action

This change adds support for goimports as a code action that can be run
on save. However, there do appear to be issues with the propagation of
the context.Only field of the CodeActionParams, so we treat every
codeAction as an organizeImports action - this should be fixed in the
next vscode-languageclient release
(https://github.com/Microsoft/vscode-languageserver-node/issues/442).

Change-Id: I64ca0034c393762248fde6521aba86ed9d41bf70
Reviewed-on: https://go-review.googlesource.com/c/154338
Reviewed-by: Ian Cottrell <iancottrell@google.com>
---
 internal/lsp/imports.go       | 34 ++++++++++++++++++++++++++++++++++
 internal/lsp/server.go        | 19 +++++++++++++++++--
 internal/lsp/source/format.go | 22 ++++++++++++++++++----
 3 files changed, 69 insertions(+), 6 deletions(-)
 create mode 100644 internal/lsp/imports.go

diff --git a/internal/lsp/imports.go b/internal/lsp/imports.go
new file mode 100644
index 00000000000..c3ecc132e13
--- /dev/null
+++ b/internal/lsp/imports.go
@@ -0,0 +1,34 @@
+// Copyright 2018 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package lsp
+
+import (
+	"context"
+
+	"golang.org/x/tools/internal/lsp/cache"
+	"golang.org/x/tools/internal/lsp/protocol"
+	"golang.org/x/tools/internal/lsp/source"
+)
+
+func organizeImports(ctx context.Context, v *cache.View, uri protocol.DocumentURI) ([]protocol.TextEdit, error) {
+	f := v.GetFile(source.URI(uri))
+	tok, err := f.GetToken()
+	if err != nil {
+		return nil, err
+	}
+	r := source.Range{
+		Start: tok.Pos(0),
+		End:   tok.Pos(tok.Size()),
+	}
+	content, err := f.Read()
+	if err != nil {
+		return nil, err
+	}
+	edits, err := source.Imports(ctx, tok.Name(), content, r)
+	if err != nil {
+		return nil, err
+	}
+	return toProtocolEdits(tok, content, edits), nil
+}
diff --git a/internal/lsp/server.go b/internal/lsp/server.go
index 21941d5c63f..cc10074930f 100644
--- a/internal/lsp/server.go
+++ b/internal/lsp/server.go
@@ -56,6 +56,7 @@ func (s *server) Initialize(ctx context.Context, params *protocol.InitializePara
 
 	return &protocol.InitializeResult{
 		Capabilities: protocol.ServerCapabilities{
+			CodeActionProvider: true,
 			CompletionProvider: protocol.CompletionOptions{
 				TriggerCharacters: []string{"."},
 			},
@@ -250,8 +251,22 @@ func (s *server) DocumentSymbol(context.Context, *protocol.DocumentSymbolParams)
 	return nil, notImplemented("DocumentSymbol")
 }
 
-func (s *server) CodeAction(context.Context, *protocol.CodeActionParams) ([]protocol.CodeAction, error) {
-	return nil, notImplemented("CodeAction")
+func (s *server) CodeAction(ctx context.Context, params *protocol.CodeActionParams) ([]protocol.CodeAction, error) {
+	edits, err := organizeImports(ctx, s.view, params.TextDocument.URI)
+	if err != nil {
+		return nil, err
+	}
+	return []protocol.CodeAction{
+		{
+			Title: "Organize Imports",
+			Kind:  protocol.SourceOrganizeImports,
+			Edit: protocol.WorkspaceEdit{
+				Changes: map[protocol.DocumentURI][]protocol.TextEdit{
+					params.TextDocument.URI: edits,
+				},
+			},
+		},
+	}, nil
 }
 
 func (s *server) CodeLens(context.Context, *protocol.CodeLensParams) ([]protocol.CodeLens, error) {
diff --git a/internal/lsp/source/format.go b/internal/lsp/source/format.go
index 2f9938f63a7..0ac932a00de 100644
--- a/internal/lsp/source/format.go
+++ b/internal/lsp/source/format.go
@@ -13,9 +13,10 @@ import (
 	"go/format"
 
 	"golang.org/x/tools/go/ast/astutil"
+	"golang.org/x/tools/imports"
 )
 
-// Format formats a document with a given range.
+// Format formats a file with a given range.
 func Format(ctx context.Context, f File, rng Range) ([]TextEdit, error) {
 	fAST, err := f.GetAST()
 	if err != nil {
@@ -54,11 +55,24 @@ func Format(ctx context.Context, f File, rng Range) ([]TextEdit, error) {
 	if err := format.Node(buf, fset, node); err != nil {
 		return nil, err
 	}
-	// TODO(rstambler): Compute text edits instead of replacing whole file.
+	return computeTextEdits(rng, buf.String()), nil
+}
+
+// Imports formats a file using the goimports tool.
+func Imports(ctx context.Context, filename string, content []byte, rng Range) ([]TextEdit, error) {
+	content, err := imports.Process(filename, content, nil)
+	if err != nil {
+		return nil, err
+	}
+	return computeTextEdits(rng, string(content)), nil
+}
+
+// TODO(rstambler): Compute text edits instead of replacing whole file.
+func computeTextEdits(rng Range, content string) []TextEdit {
 	return []TextEdit{
 		{
 			Range:   rng,
-			NewText: buf.String(),
+			NewText: content,
 		},
-	}, nil
+	}
 }