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>
diff --git a/internal/lsp/code_action.go b/internal/lsp/code_action.go
index 425bfe3..b875798 100644
--- a/internal/lsp/code_action.go
+++ b/internal/lsp/code_action.go
@@ -432,7 +432,7 @@
 			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 6c16c52..2695095 100644
--- a/internal/lsp/lsp_test.go
+++ b/internal/lsp/lsp_test.go
@@ -454,12 +454,12 @@
 		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 @@
 	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 b6b9f7d..4725572 100644
--- a/internal/lsp/source/extract.go
+++ b/internal/lsp/source/extract.go
@@ -45,32 +45,32 @@
 		}
 		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