From 5c72ddda614e1a1d966e7dae667c593cdbbdac4b Mon Sep 17 00:00:00 2001
From: Josh Baum <joshbaum@google.com>
Date: Mon, 3 Aug 2020 12:16:31 -0400
Subject: [PATCH] internal/lsp: fix bug in extract variable edit positions

Previously, the suggested fix tests did not properly handle the case
in which one fix contained at least two edits. We also prevent
the server from panicing when we cannot extract the selection.

Change-Id: I38f7b6d871b2f2741349a3fd94fd95b396f5fd33
Reviewed-on: https://go-review.googlesource.com/c/tools/+/246458
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
---
 internal/lsp/code_action.go    |  2 +-
 internal/lsp/lsp_test.go       | 34 ++++++++++++++++++++++++++++++++--
 internal/lsp/source/extract.go | 20 ++++++++++----------
 3 files changed, 43 insertions(+), 13 deletions(-)

diff --git a/internal/lsp/code_action.go b/internal/lsp/code_action.go
index 425bfe35db..b8757986c8 100644
--- a/internal/lsp/code_action.go
+++ b/internal/lsp/code_action.go
@@ -432,7 +432,7 @@ func extractionFixes(ctx context.Context, snapshot source.Snapshot, pkg source.P
 			Title: command.Title,
 			Kind:  protocol.RefactorExtract,
 			Command: &protocol.Command{
-				Command:   source.CommandExtractFunction.Name,
+				Command:   command.Name,
 				Arguments: jsonArgs,
 			},
 		})
diff --git a/internal/lsp/lsp_test.go b/internal/lsp/lsp_test.go
index 6c16c520e5..2695095071 100644
--- a/internal/lsp/lsp_test.go
+++ b/internal/lsp/lsp_test.go
@@ -454,12 +454,12 @@ func (r *runner) SuggestedFix(t *testing.T, spn span.Span, actionKinds []string)
 		if err != nil {
 			t.Fatal(err)
 		}
-		res, err = applyTextDocumentEdits(r, edits)
+		res, err = applySuggestedFixEdits(r, edits)
 		if err != nil {
 			t.Fatal(err)
 		}
 	} else {
-		res, err = applyTextDocumentEdits(r, action.Edit.DocumentChanges)
+		res, err = applySuggestedFixEdits(r, action.Edit.DocumentChanges)
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -872,6 +872,36 @@ func applyTextDocumentEdits(r *runner, edits []protocol.TextDocumentEdit) (map[s
 	return res, nil
 }
 
+func applySuggestedFixEdits(r *runner, edits []protocol.TextDocumentEdit) (map[span.URI]string, error) {
+	res := map[span.URI]string{}
+	for _, docEdits := range edits {
+		uri := docEdits.TextDocument.URI.SpanURI()
+		var m *protocol.ColumnMapper
+		// If we have already edited this file, we use the edited version (rather than the
+		// file in its original state) so that we preserve our initial changes.
+		if content, ok := res[uri]; ok {
+			m = &protocol.ColumnMapper{
+				URI: uri,
+				Converter: span.NewContentConverter(
+					uri.Filename(), []byte(content)),
+				Content: []byte(content),
+			}
+		} else {
+			var err error
+			if m, err = r.data.Mapper(uri); 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/source/extract.go b/internal/lsp/source/extract.go
index b6b9f7da3f..4725572295 100644
--- a/internal/lsp/source/extract.go
+++ b/internal/lsp/source/extract.go
@@ -45,32 +45,32 @@ func extractVariable(fset *token.FileSet, rng span.Range, src []byte, file *ast.
 		}
 		assignment = buf.String()
 	case *ast.CallExpr: // TODO: find number of return values and do according actions.
-		return nil, nil
+		return nil, fmt.Errorf("cannot extract call expression")
 	default:
-		return nil, nil
+		return nil, fmt.Errorf("cannot extract %T", expr)
 	}
 
 	insertBeforeStmt := analysisinternal.StmtToInsertVarBefore(path)
 	if insertBeforeStmt == nil {
-		return nil, nil
+		return nil, fmt.Errorf("cannot find location to insert extraction")
 	}
 
 	tok := fset.File(expr.Pos())
 	if tok == nil {
-		return nil, nil
+		return nil, fmt.Errorf("no file for pos %v", fset.Position(file.Pos()))
 	}
 	indent := calculateIndentation(src, tok, insertBeforeStmt)
 	return &analysis.SuggestedFix{
 		TextEdits: []analysis.TextEdit{
 			{
-				Pos:     insertBeforeStmt.Pos(),
-				End:     insertBeforeStmt.End(),
-				NewText: []byte(assignment + "\n" + indent),
+				Pos:     rng.Start,
+				End:     rng.End,
+				NewText: []byte(name),
 			},
 			{
-				Pos:     rng.Start,
-				End:     rng.Start,
-				NewText: []byte(name),
+				Pos:     insertBeforeStmt.Pos(),
+				End:     insertBeforeStmt.Pos(),
+				NewText: []byte(assignment + "\n" + indent),
 			},
 		},
 	}, nil