sorbet · neilparikh · Jan 29, 2025 · Feb 6, 2025 · Feb 6, 2025 · Apr 30, 2025
@@ -115,6 +115,7 @@ void Literal::_sanityCheck() {
         case core::TypePtr::Tag::FloatLiteralType:
         case core::TypePtr::Tag::NamedLiteralType:
         case core::TypePtr::Tag::ClassType:
+            // If you add a new valid tag here, make sure to also update the code in TreeEquality.cc to handle that tag
             break;
         default:
             ENFORCE(false, "unexpected TypePtr::Tag: {}", fmt::underlying(tag));

@@ -31,6 +31,7 @@ cc_library(
     ],
     hdrs = [
         "ConvertToSingletonClassMethod.h",
+        "ExtractMethod.h",
         "ExtractVariable.h",
         "LSPConfiguration.h",
         "LSPInput.h",

@@ -0,0 +1,214 @@
+#include "main/lsp/ExtractMethod.h"
+#include "absl/strings/str_join.h"
+#include "ast/treemap/treemap.h"
+
+using namespace std;
+
+namespace sorbet::realmain::lsp {
+
+class Walk {
+    // The selection loc
+    core::LocOffsets targetLoc;
+
+public:
+    // After the walk is complete, this should point to the deepest insSeq that contains targetLoc
+    const ast::InsSeq *enclosingInsSeq;
+    const ast::MethodDef *enclosingMethodDef;
+    Walk(core::LocOffsets targetLoc) : targetLoc(targetLoc), enclosingInsSeq(nullptr) {}
+
+    void preTransformInsSeq(core::Context ctx, const ast::ExpressionPtr &tree) {
+        auto &insSeq = ast::cast_tree_nonnull<ast::InsSeq>(tree);
+        if (insSeq.loc.exists() && insSeq.loc.contains(targetLoc)) {
+            if (!enclosingInsSeq || enclosingInsSeq->loc.contains(insSeq.loc)) {
+                enclosingInsSeq = &insSeq;
+            }
+        }
+    }
+
+    void preTransformMethodDef(core::Context ctx, const ast::ExpressionPtr &tree) {
+        auto &methodDef = ast::cast_tree_nonnull<ast::MethodDef>(tree);
+        if (methodDef.loc.exists() && methodDef.loc.contains(targetLoc)) {
+            enclosingMethodDef = &methodDef;
+        }
+    }
+};
+
+class FindLocalsWalk {
+    absl::flat_hash_set<core::NameRef> &localsAssigned;
+    absl::flat_hash_set<core::NameRef> &localsReferenced;
+
+    absl
10000
::flat_hash_set<core::LocOffsets> lhsLocs;
+
+public:
+    FindLocalsWalk(absl::flat_hash_set<core::NameRef> &localsAssigned,
+                   absl::flat_hash_set<core::NameRef> &localsReferenced)
+        : localsAssigned(localsAssigned), localsReferenced(localsReferenced) {}
+    void preTransformAssign(core::Context ctx, const ast::Assign &assign) {
+        if (auto local = ast::cast_tree<ast::Local>(assign.lhs)) {
+            localsAssigned.insert(local->localVariable._name);
+            lhsLocs.insert(local->loc);
+        }
+        return;
+    }
+
+    void postTransformLocal(core::Context ctx, const ast::Local &local) {
+        if (!lhsLocs.contains(local.loc)) {
+            localsReferenced.insert(local.localVariable._name);
+        }
+    }
+};
+
+void findLocals(const core::GlobalState &gs, const core::FileRef file,
+                absl::flat_hash_set<core::NameRef> &localsAssigned,
+                absl::flat_hash_set<core::NameRef> &localsReferenced, const ast::ExpressionPtr *insn) {
+    FindLocalsWalk walk(localsAssigned, localsReferenced);
+    core::Context ctx(gs, core::Symbols::root(), file);
+    ast::ConstTreeWalk::apply(ctx, walk, *insn);
+}
+
+const ast::ExpressionPtr *getElem(const ast::InsSeq *insSeq, int i) {
+    if (i > insSeq->stats.size()) {
+        return nullptr;
+    }
+
+    if (i == insSeq->stats.size()) {
+        return &insSeq->expr;
+    }
+
+    return &insSeq->stats[i];
+}
+
+vector<unique_ptr<TextDocumentEdit>> MethodExtractor::getExtractEdits(const LSPTypecheckerDelegate &typechecker,
+                                                                      const LSPConfiguration &config) {
+    const auto file = selectionLoc.file();
+    const auto &gs = typechecker.state();
+    // Steps:
+    // - Find enclosing InsSeq (insSeq.contains(selectionLoc))
+    Walk walk(selectionLoc.offsets());
+    // TODO(neil): do we need local Vars?
+    auto afterLocalVars = typechecker.getLocalVarTrees(file);
+    core::Context ctx(gs, core::Symbols::root(), file);
+    ast::TreeWalk::apply(ctx, walk, afterLocalVars);
+    auto *enclosingInsSeq = walk.enclosingInsSeq;
+    auto *enclosingMethodDef = walk.enclosingMethodDef;
+
+    if (!enclosingInsSeq || !enclosingMethodDef) {
+        return {};
+    }
+    //
+    // - Step through each insn in InsSeq, and see if selection corrosponds to full insns
+    //   - start of selection == start of some insn and end of selection == end of some insn
+    int start_idx = -1;
+    int end_idx = -1;
+    // TODO: use vector instead of set so it's ordered? or btree_set?
+    absl::flat_hash_set<core::NameRef> localsAssignedBefore;
+    // TODO: replace with std::transform/absl::c_transform?
+    for (auto &param : enclosingMethodDef->args) {
+        if (auto local = ast::cast_tree<ast::Local>(param)) {
+            localsAssignedBefore.insert(local->localVariable._name);
+        }
+    }
+    absl::flat_hash_set<core::NameRef> localsReferencedIn;
+    absl::flat_hash_set<core::NameRef> localsAssignedIn;
+    absl::flat_hash_set<core::NameRef> localsReferencedAfter;
+    // size() + 1 to account for InsSeq#expr
+    // TODO: we need to find all the locals from before the enclosing insSeq too
+    for (int i = 0; i < enclosingInsSeq->stats.size() + 1; i++) {
+        auto *insn = getElem(enclosingInsSeq, i);
+        if (start_idx == -1) {
+            // TODO: move the start of selectionLoc to the first non-whitespace character, so we can ignore the
+            // indentation at the start
+            if (insn->loc().beginPos() == selectionLoc.offsets().beginPos()) {
+                // First insn in the user's selection
+                start_idx = i;
+
+                findLocals(gs, file, localsAssignedIn, localsReferencedIn, insn);
+            } else {
+                // Before the user's selection
+                absl::flat_hash_set<core::NameRef> localsReferencedBefore_;
+                findLocals(gs, file, localsAssignedBefore, localsReferencedBefore_, insn);
+            }
+        } else if (end_idx == -1) {
+            // Inside the user's selection
+            findLocals(gs, file, localsAssignedIn, localsReferencedIn, insn);
+
+            // TODO: move the end of selectionLoc to the last non-whitespace character, so we can ignore the
+            // indentation at the end
+            if (insn->loc().endPos() == selectionLoc.offsets().endPos()) {
+                // Last insn in the user's selection
+                end_idx = i;
+            }
+        } else {
+            // After the user's selection
+            absl::flat_hash_set<core::NameRef> localsAssignedAfter_;
+            findLocals(gs, file, localsAssignedAfter_, localsReferencedAfter, insn);
+        }
+    }
+
+    if (start_idx == -1 || end_idx == -1) {
+        return {};
+    }
+    /* for (auto name : localsAssignedBefore) { */
+    /*     fmt::print("assignedBefore: {}\n", name.toString(gs)); */
+    /* } */
+    /* for (auto name : localsAssignedIn) { */
+    /*     fmt::print("assignedIn: {}\n", name.toString(gs)); */
+    /* } */
+    /* for (auto name : localsReferencedIn) { */
+    /*     fmt::print("referencedIn: {}\n", name.toString(gs)); */
+    /* } */
+    /* for (auto name : localsReferencedAfter) { */
+    /*     fmt::print("referencedAfter: {}\n", name.toString(gs)); */
+    /* } */
+    // TODO: if the selection includes a return, don't allow extracting it for now?
+
+    std::vector<string> methodParams;
+    std::vector<string> methodReturns;
+    // args for new method are those that are (referencedIn and (assignedBefore or !assignedIn))
+    for (auto name : localsReferencedIn) {
+        if (name == core::Names::selfLocal()) {
+            continue;
+        }
+        if (localsAssignedBefore.contains(name) || !localsAssignedIn.contains(name)) {
+            methodParams.push_back(name.show(gs));
+        }
+    }
+    // new method will return those are (assignedIn and referencedAfter)
+    for (auto name : localsAssignedIn) {
+        /* if (localsReferencedAfter.contains(name)) { */
+        methodReturns.push_back(name.show(gs));
+        /* } */
+    }
+    // TODO: sort methodParams and methodReturns for deterministic behaviour
+
+    string firstLine = fmt::format("def newMethod({})", absl::StrJoin(methodParams, ", "));
+    string lastLine = fmt::format("return {}", absl::StrJoin(methodReturns, ", "));
+    auto [_, numSpaces] = selectionLoc.findStartOfIndentation(gs);
+    config.logger->info("{}\n", numSpaces);
+    auto bodyIndentation = std::string(numSpaces, ' ');
+    auto defEndIndentation = std::string(numSpaces - 2, ' ');
+    auto newMethod =
+        fmt::format("{}\n{}{}\n{}{}\n{}end\n\n{}", firstLine, bodyIndentation, selectionLoc.source(gs).value(),
+                    bodyIndentation, lastLine, defEndIndentation, defEndIndentation);
+
+    vector<unique_ptr<TextEdit>> edits;
+    edits.emplace_back(make_unique<TextEdit>(
+        Range::fromLoc(gs, core::Loc(file, enclosingMethodDef->loc.copyWithZeroLength())), newMethod));
+    if (methodReturns.size() > 0) {
+        edits.emplace_back(make_unique<TextEdit>(
+            Range::fromLoc(gs, selectionLoc),
+            fmt::format("{} = newMethod({})", absl::StrJoin(methodReturns, ", "), absl::StrJoin(methodParams, ", "))));
+    } else {
+        edits.emplace_back(make_unique<TextEdit>(Range::fromLoc(gs, selectionLoc),
+                                                 fmt::format("newMethod({})", absl::StrJoin(methodParams, ", "))));
+    }
+
+    auto docEdit = make_unique<TextDocumentEdit>(
+        make_unique<VersionedTextDocumentIdentifier>(config.fileRef2Uri(gs, file), JSONNullObject()), move(edits));
+
+    vector<unique_ptr<TextDocumentEdit>> res;
+    res.emplace_back(move(docEdit));
+    return res;
+}
+
+} // namespace sorbet::realmain::lsp
@@ -0,0 +1,22 @@
+#ifndef SORBET_EXTRACT_METHOD_H
+#define SORBET_EXTRACT_METHOD_H
+
+#include "main/lsp/LSPConfiguration.h"
+#include "main/lsp/LSPTypechecker.h"
+#include "main/lsp/json_types.h"
+
+namespace sorbet::realmain::lsp {
+
+class MethodExtractor {
+    const core::Loc selectionLoc;
+    ast::ExpressionPtr matchingNode;
+
+public:
+    MethodExtractor(const core::Loc selectionLoc) : selectionLoc(selectionLoc) {}
+    std::vector<std::unique_ptr<TextDocumentEdit>> getExtractEdits(const LSPTypecheckerDelegate &typechecker,
+                                                                   const LSPConfiguration &config);
+};
+
+} // namespace sorbet::realmain::lsp
+
+#endif
@@ -106,6 +106,7 @@ core::LocOffsets findWhereToInsert(const ast::ExpressionPtr &scope, const core::
 // This tree walk takes a Loc and looks for nodes that have that Loc exactly
 class LocSearchWalk {
     // The selection loc
+    // TODO(neil): change this to LocOffsets
     core::Loc targetLoc;
     // At the end of this walk, we want to return what class/method the matching expression was part of.
     // To do that, we can maintain a stack of classes/method, so that when we get a match, we can capture

@@ -4,6 +4,7 @@
 #include "common/sort/sort.h"
 #include "core/lsp/QueryResponse.h"
 #include "main/lsp/ConvertToSingletonClassMethod.h"
+#include "main/lsp/ExtractMethod.h"
 #include "main/lsp/ExtractVariable.h"
 #include "main/lsp/LSPLoop.h"
 #include "main/lsp/LSPQuery.h"
@@ -314,6 +315,24 @@ unique_ptr<ResponseMessage> CodeActionTask::runRequest(LSPTypecheckerDelegate &t
                 // TODO(neil): trigger a rename for newVariable
             }
         }
+        if (config.opts.lspExtractToMethodEnabled) {
+            MethodExtractor methodExtractor(loc);
+            vector<unique_ptr<TextDocumentEdit>> documentEdits;
+            {
+                Timer timeit(gs.tracer(), "Extract to Method");
+                documentEdits = methodExtractor.getExtractEdits(typechecker, config);
+            }
+            if (!documentEdits.empty()) {
+                auto action = make_unique<CodeAction>("Extract Method");
+                action->kind = CodeActionKind::RefactorExtract;
+
+                auto workspaceEdit = make_unique<WorkspaceEdit>();
+                workspaceEdit->documentChanges = move(documentEdits);
+
+                action->edit = move(workspaceEdit);
+                result.emplace_back(move(action));
+            }
+        }
     }
 
     response->result = move(result);

@@ -549,6 +549,8 @@ buildOptions(const vector<pipeline::semantic_extension::SemanticExtensionProvide
 
     options.add_options(section)("enable-experimental-lsp-extract-to-variable",
                                  "Enable experimental LSP feature: Extract To Variable");
+    options.add_options(section)("enable-experimental-lsp-extract-to-method",
+                                 "Enable experimental LSP feature: Extract To Method");
     options.add_options(section)(
         "enable-all-experimental-lsp-features",
         "Enable every experimental LSP feature. (WARNING: can be crashy; for developer use only. "
@@ -932,6 +934,8 @@ void readOptions(Options &opts,
         opts.lspAllBetaFeaturesEnabled = enableAllLSPFeatures || raw["enable-all-beta-lsp-features"].as<bool>();
         opts.lspExtractToVariableEnabled =
             opts.lspAllBetaFeaturesEnabled || raw["enable-experimental-lsp-extract-to-variable"].as<bool>();
+        opts.lspExtractToMethodEnabled =
+            opts.lspAllBetaFeaturesEnabled || raw["enable-experimental-lsp-extract-to-method"].as<bool>();
         opts.lspDocumentHighlightEnabled =
             enableAllLSPFeatures || raw["enable-experimental-lsp-document-highlight"].as<bool>();
         opts.lspSignatureHelpEnabled = enableAllLSPFeatures || raw["enable-experimental-lsp-signature-help"].as<bool>();

@@ -283,6 +283,7 @@ struct Options {
     bool lspDocumentFormatRubyfmtEnabled = false;
     bool lspSignatureHelpEnabled = false;
     bool lspExtractToVariableEnabled = false;
+    bool lspExtractToMethodEnabled = false;
     bool forciblySilenceLspMultipleDirError = false;
     // Enables out-of-order reference checking
     bool outOfOrderReferenceChecksEnabled = false;

@@ -261,6 +261,7 @@ const UnorderedMap<
         {"typed-super", BooleanPropertyAssertion::make},
         {"enable-suggest-unsafe", BooleanPropertyAssertion::make},
         {"enable-experimental-lsp-extract-to-variable", BooleanPropertyAssertion::make},
+        {"enable-experimental-lsp-extract-to-method", BooleanPropertyAssertion::make},
            {"selective-apply-code-action", StringPropertyAssertions::make},
         {"use-code-action-resolve", BooleanPropertyAssertion::make},
         {"assert-no-code-action", StringPropertyAssertions::make},

@@ -548,6 +548,8 @@ TEST_CASE("LSPTest") {
         opts.lspExtractToVariableEnabled =
             BooleanPropertyAssertion::getValue("enable-experimental-lsp-extract-to-variable", assertions)
                 .value_or(false);
+        opts.lspExtractToMethodEnabled =
+            BooleanPropertyAssertion::getValue("enable-experimental-lsp-extract-to-method", assertions).value_or(false);
         opts.disableWatchman = true;
         opts.rubyfmtPath = "test/testdata/lsp/rubyfmt-stub/rubyfmt";
 

@@ -0,0 +1,19 @@
+# typed: true
+# selective-apply-code-action: refactor.extract
+# enable-experimental-lsp-extract-to-method: true
+
+class A
+  def y
+    1
+  end
+
+  def newMethod(x, a)
+    b = x * y; c = x + y + a + b
+    return c, b
+  end
+
+  def foo(x)
+    a = x + y; c, b = newMethod(x, a); return c
+  #            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ apply-code-action: [A] Extract Method
+  end
+end
@@ -0,0 +1,14 @@
+# typed: true
+# selective-apply-code-action: refactor.extract
+# enable-experimental-lsp-extract-to-method: true
+
+class A
+  def y
+    1
+  end
+
+  def foo(x)
+    a = x + y; b = x * y; c = x + y + a + b; return c
  #            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ apply-code-action: [A] Extract Method
+  end
+end
@@ -244,6 +244,8 @@ Usage:
                                 Enable experimental LSP feature: Signature Help
       --enable-experimental-lsp-extract-to-variable
                                 Enable experimental LSP feature: Extract To Variable
+      --enable-experimental-lsp-extract-to-method
+                                Enable experimental LSP feature: Extract To Method
       --enable-all-experimental-lsp-features
                                 Enable every experimental LSP feature. (WARNING: can be
                                 crashy; for developer use only. End users should prefer