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

@@ -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,
 			},
 		})

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
 

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