internal/lsp: support extract function

Extract function is a code action, similar to extract variable. After highlighting a selection, if valid, the lightbulb appears to trigger extraction. The current implementation does not allow users to extract selections with a return statement. Updates golang/go#37170 Change-Id: I5fc3b19cf7dbca4407ecf0cc37017661223614d1 Reviewed-on: https://go-review.googlesource.com/c/tools/+/241957 Run-TryBot: Rebecca Stambler <[email protected]> Run-TryBot: Josh Baum <[email protected]> TryBot-Result: Gobot Gobot <[email protected]> Reviewed-by: Rebecca Stambler <[email protected]>
golang · Jul 20, 2020 · 6d307ed · 6d307ed
1 parent 9cbb971
commit 6d307ed
Show file tree

Hide file tree

Showing 27 changed files with 844 additions and 34 deletions.
diff --git a/internal/analysisinternal/analysis.go b/internal/analysisinternal/analysis.go
@@ -84,9 +84,95 @@ func TypeExpr(fset *token.FileSet, f *ast.File, pkg *types.Package, typ types.Ty
 		default:
 			return ast.NewIdent(t.Name())
 		}
+	case *types.Pointer:
+		x := TypeExpr(fset, f, pkg, t.Elem())
+		if x == nil {
+			return nil
+		}
+		return &ast.UnaryExpr{
+			Op: token.MUL,
+			X:  x,
+		}
+	case *types.Array:
+		elt := TypeExpr(fset, f, pkg, t.Elem())
+		if elt == nil {
+			return nil
+		}
+		return &ast.ArrayType{
+			Len: &ast.BasicLit{
+				Kind:  token.INT,
+				Value: fmt.Sprintf("%d", t.Len()),
+			},
+			Elt: elt,
+		}
+	case *types.Slice:
+		elt := TypeExpr(fset, f, pkg, t.Elem())
+		if elt == nil {
+			return nil
+		}
+		return &ast.ArrayType{
+			Elt: elt,
+		}
+	case *types.Map:
+		key := TypeExpr(fset, f, pkg, t.Key())
+		value := TypeExpr(fset, f, pkg, t.Elem())
+		if key == nil || value == nil {
+			return nil
+		}
+		return &ast.MapType{
+			Key:   key,
+			Value: value,
+		}
+	case *types.Chan:
+		dir := ast.ChanDir(t.Dir())
+		if t.Dir() == types.SendRecv {
+			dir = ast.SEND | ast.RECV
+		}
+		value := TypeExpr(fset, f, pkg, t.Elem())
+		if value == nil {
+			return nil
+		}
+		return &ast.ChanType{
+			Dir:   dir,
+			Value: value,
+		}
+	case *types.Signature:
+		var params []*ast.Field
+		for i := 0; i < t.Params().Len(); i++ {
+			p := TypeExpr(fset, f, pkg, t.Params().At(i).Type())
+			if p == nil {
+				return nil
+			}
+			params = append(params, &ast.Field{
+				Type: p,
+				Names: []*ast.Ident{
+					{
+						Name: t.Params().At(i).Name(),
+					},
+				},
+			})
+		}
+		var returns []*ast.Field
+		for i := 0; i < t.Results().Len(); i++ {
+			r := TypeExpr(fset, f, pkg, t.Results().At(i).Type())
+			if r == nil {
+				return nil
+			}
+			returns = append(returns, &ast.Field{
+				Type: r,
+			})
+		}
+		return &ast.FuncType{
+			Params: &ast.FieldList{
+				List: params,
+			},
+			Results: &ast.FieldList{
+				List: returns,
+			},
+		}
 	case *types.Named:
 		if t.Obj().Pkg() == nil {
-			return nil
+			return ast.NewIdent(t.Obj().Name())
 		}
 		if t.Obj().Pkg() == pkg {
 			return ast.NewIdent(t.Obj().Name())
@@ -109,11 +195,6 @@ func TypeExpr(fset *token.FileSet, f *ast.File, pkg *types.Package, typ types.Ty
 			X:   ast.NewIdent(pkgName),
 			Sel: ast.NewIdent(t.Obj().Name()),
 		}
-	case *types.Pointer:
-		return &ast.UnaryExpr{
-			Op: token.MUL,
-			X:  TypeExpr(fset, f, pkg, t.Elem()),
-		}
 	default:
 		return nil // TODO: anonymous structs, but who does that
 	}

diff --git a/internal/lsp/cmd/test/cmdtest.go b/internal/lsp/cmd/test/cmdtest.go
@@ -133,6 +133,10 @@ func (r *runner) RankCompletion(t *testing.T, src span.Span, test tests.Completi
 	//TODO: add command line completions tests when it works
 }
 
+func (r *runner) FunctionExtraction(t *testing.T, start span.Span, end span.Span) {
+	//TODO: function extraction not supported on command line
+}
+
 func (r *runner) runGoplsCmd(t testing.TB, args ...string) (string, string) {
 	rStdout, wStdout, err := os.Pipe()
 	if err != nil {

diff --git a/internal/lsp/code_action.go b/internal/lsp/code_action.go
@@ -361,22 +361,34 @@ func extractionFixes(ctx context.Context, snapshot source.Snapshot, ph source.Pa
 	if err != nil {
 		return nil, nil
 	}
+	var actions []protocol.CodeAction
 	edits, err := source.ExtractVariable(ctx, snapshot, fh, rng)
 	if err != nil {
 		return nil, err
 	}
-	if len(edits) == 0 {
-		return nil, nil
-	}
-	return []protocol.CodeAction{
-		{
+	if len(edits) > 0 {
+		actions = append(actions, protocol.CodeAction{
 			Title: "Extract to variable",
 			Kind:  protocol.RefactorExtract,
 			Edit: protocol.WorkspaceEdit{
 				DocumentChanges: documentChanges(fh, edits),
 			},
-		},
-	}, nil
+		})
+	}
+	edits, err = source.ExtractFunction(ctx, snapshot, fh, rng)
+	if err != nil {
+		return nil, err
+	}
+	if len(edits) > 0 {
+		actions = append(actions, protocol.CodeAction{
+			Title: "Extract to function",
+			Kind:  protocol.RefactorExtract,
+			Edit: protocol.WorkspaceEdit{
+				DocumentChanges: documentChanges(fh, edits),
+			},
+		})
+	}
+	return actions, nil
 }
 
 func documentChanges(fh source.FileHandle, edits []protocol.TextEdit) []protocol.TextDocumentEdit {

diff --git a/internal/lsp/lsp_test.go b/internal/lsp/lsp_test.go
@@ -433,6 +433,52 @@ func (r *runner) SuggestedFix(t *testing.T, spn span.Span, actionKinds []string)
 	}
 }
 
+func (r *runner) FunctionExtraction(t *testing.T, start span.Span, end span.Span) {
+	uri := start.URI()
+	_, err := r.server.session.ViewOf(uri)
+	if err != nil {
+		t.Fatal(err)
+	}
+	m, err := r.data.Mapper(uri)
+	if err != nil {
+		t.Fatal(err)
+	}
+	spn := span.New(start.URI(), start.Start(), end.End())
+	rng, err := m.Range(spn)
+	if err != nil {
+		t.Fatal(err)
+	}
+	actions, err := r.server.CodeAction(r.ctx, &protocol.CodeActionParams{
+		TextDocument: protocol.TextDocumentIdentifier{
+			URI: protocol.URIFromSpanURI(uri),
+		},
+		Range: rng,
+		Context: protocol.CodeActionContext{
+			Only: []protocol.CodeActionKind{"refactor.extract"},
+		},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	// Hack: We assume that we only get one code action per range.
+	// TODO(rstambler): Support multiple code actions per test.
+	if len(actions) == 0 || len(actions) > 1 {
+		t.Fatalf("unexpected number of code actions, want 1, got %v", len(actions))
+	}
+	res, err := applyWorkspaceEdits(r, actions[0].Edit)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for u, got := range res {
+		want := string(r.data.Golden("functionextraction_"+tests.SpanName(spn), u.Filename(), func() ([]byte, error) {
+			return []byte(got), nil
+		}))
+		if want != got {
+			t.Errorf("function extraction failed for %s:\n%s", u.Filename(), tests.Diff(want, got))
+		}
+	}
+}
+
 func (r *runner) Definition(t *testing.T, spn span.Span, d tests.Definition) {
 	sm, err := r.data.Mapper(d.Src.URI())
 	if err != nil {