Detach the "Move method" logic from the LSPTask (sorbet#5537)

* Move `queryBy*` functions from the `LSPTask` to a separate file

* Move the `getRenameEdits` function from the `LSPTask` to the `AbstractRenamer`

* Move the refactoring logic from the code action into separate file

* Group query functions into a LSPQuery class with static methods

* main/lsp/requests/move_method -> main/lsp/MoveMethod

* Formatting

main/lsp/AbstractRenamer.cc

@@ -1,5 +1,7 @@
 #include "AbstractRenamer.h"
+#include "main/lsp/LSPQuery.h"
 #include "main/lsp/lsp.h"
+
 using namespace std;
 
 namespace sorbet::realmain::lsp {
@@ -132,4 +134,34 @@ void AbstractRenamer::addDispatchRelatedMethods(const core::GlobalState &gs, con
     }
 }
 
+void AbstractRenamer::getRenameEdits(LSPTypecheckerInterface &typechecker, core::SymbolRef symbol, string newName) {
+    const core::GlobalState &gs = typechecker.state();
+    auto originalName = symbol.name(gs).show(gs);
+
+    addSymbol(symbol);
+
+    auto symbolQueue = getQueue();
+    for (auto sym = symbolQueue->pop(); sym.exists(); sym = symbolQueue->pop()) {
+        auto queryResult = LSPQuery::bySymbol(config, typechecker, sym);
+        if (queryResult.error) {
+            return;
+        }
+
+        // Filter for untyped files, and deduplicate responses by location.  We don't use extractLocations here because
+        // in some cases like sends, we need the SendResponse to be able to accurately find the method name in the
+        // expression.
+        for (auto &response : LSPQuery::filterAndDedup(gs, queryResult.responses)) {
+            auto loc = response->getLoc();
+            if (loc.file().data(gs).isPayload()) {
+                // We don't support renaming things in payload files.
+                return;
+            }
+
+            // We may process the same send multiple times in case of union types, but this is ok because the renamer
+            // de-duplicates edits at the same location
+            rename(response, sym);
+        }
+    }
+}
+
 } // namespace sorbet::realmain::lsp

main/lsp/AbstractRenamer.h

@@ -3,6 +3,7 @@
 
 #include "core/lsp/QueryResponse.h"
 #include "main/lsp/LSPConfiguration.h"
+#include "main/lsp/LSPTypechecker.h"
 #include "main/lsp/json_types.h"
 
 namespace sorbet::realmain::lsp {
@@ -29,6 +30,8 @@ class AbstractRenamer {
     std::variant<JSONNullObject, std::unique_ptr<WorkspaceEdit>> buildWorkspaceEdit();
     virtual void addSymbol(const core::SymbolRef) = 0;
 
+    void getRenameEdits(LSPTypecheckerInterface &typechecker, core::SymbolRef symbol, std::string newName);
+
     bool getInvalid();
     std::string getError();
     std::shared_ptr<UniqueSymbolQueue> getQueue();

main/lsp/BUILD

@@ -38,6 +38,8 @@ cc_library(
         "LSPMessage.h",
         "LSPOutput.h",
         "LSPPreprocessor.h",
+        "LSPQuery.h",
+        "MoveMethod.h",
         "json_types.h",
         "lsp.h",
         "wrapper.h",

main/lsp/LSPQuery.cc

@@ -0,0 +1,124 @@
+#include "main/lsp/LSPQuery.h"
+#include "common/sort.h"
+#include "core/NameHash.h"
+#include "main/lsp/json_types.h"
+
+using namespace std;
+
+namespace sorbet::realmain::lsp {
+
+// Filter for untyped locations, and dedup responses that are at the same location
+vector<unique_ptr<core::lsp::QueryResponse>>
+LSPQuery::filterAndDedup(const core::GlobalState &gs,
+                         const vector<unique_ptr<core::lsp::QueryResponse>> &queryResponses) {
+    vector<unique_ptr<core::lsp::QueryResponse>> responses;
+    // Filter for responses with a loc that exists and points to a typed file, unless it's a const, field or
+    // definition in which case we're ok with untyped files (because we know where those things are even in untyped
+    // files.)
+    for (auto &q : queryResponses) {
+        core::Loc loc = q->getLoc();
+        if (loc.exists() && loc.file().exists()) {
+            auto fileIsTyped = loc.file().data(gs).strictLevel >= core::StrictLevel::True;
+            // If file is untyped, only support responses involving constants and definitions.
+            if (fileIsTyped || q->isConstant() || q->isField() || q->isMethodDef()) {
+                responses.push_back(make_unique<core::lsp::QueryResponse>(*q));
+            }
+        }
+    }
+
+    // sort by location and deduplicate
+    fast_sort(responses,
+              [](const unique_ptr<core::lsp::QueryResponse> &a, const unique_ptr<core::lsp::QueryResponse> &b) -> bool {
+                  auto aLoc = a->getLoc();
+                  auto bLoc = b->getLoc();
+                  int cmp = aLoc.file().id() - bLoc.file().id();
+                  if (cmp == 0) {
+                      cmp = aLoc.beginPos() - bLoc.beginPos();
+                  }
+                  if (cmp == 0) {
+                      cmp = aLoc.endPos() - bLoc.endPos();
+                  }
+                  // TODO: precedence based on response type, in case of same location?
+                  return cmp < 0;
+              });
+    responses.resize(
+        std::distance(responses.begin(), std::unique(responses.begin(), responses.end(),
+                                                     [](const unique_ptr<core::lsp::QueryResponse> &a,
+                                                        const unique_ptr<core::lsp::QueryResponse> &b) -> bool {
+                                                         auto aLoc = a->getLoc();
+                                                         auto bLoc = b->getLoc();
+                                                         return aLoc == bLoc;
+                                                     })));
+    return responses;
+}
+
+LSPQueryResult LSPQuery::byLoc(const LSPConfiguration &config, LSPTypecheckerInterface &typechecker, string_view uri,
+                               const Position &pos, LSPMethod forMethod, bool errorIfFileIsUntyped) {
+    Timer timeit(config.logger, "setupLSPQueryByLoc");
+    const core::GlobalState &gs = typechecker.state();
+    auto fref = config.uri2FileRef(gs, uri);
+
+    if (!fref.exists() && config.isFileIgnored(config.remoteName2Local(uri))) {
+        auto error = make_unique<ResponseError>(
+            (int)LSPErrorCodes::InvalidParams,
+            fmt::format("Ignored file at uri {} in {}", uri, convertLSPMethodToString(forMethod)));
+        return LSPQueryResult{{}, move(error)};
+    }
+
+    if (!fref.exists()) {
+        auto error = make_unique<ResponseError>(
+            (int)LSPErrorCodes::InvalidParams,
+            fmt::format("Did not find file at uri {} in {}", uri, convertLSPMethodToString(forMethod)));
+        return LSPQueryResult{{}, move(error)};
+    }
+
+    if (errorIfFileIsUntyped && fref.data(gs).strictLevel < core::StrictLevel::True) {
+        config.logger->info("Ignoring request on untyped file `{}`", uri);
+        // Act as if the query returned no results.
+        return LSPQueryResult{{}, nullptr};
+    }
+
+    auto loc = config.lspPos2Loc(fref, pos, gs);
+    return typechecker.query(core::lsp::Query::createLocQuery(loc), {fref});
+}
+
+LSPQueryResult LSPQuery::LSPQuery::bySymbolInFiles(const LSPConfiguration &config, LSPTypecheckerInterface &typechecker,
+                                                   core::SymbolRef symbol, vector<core::FileRef> frefs) {
+    Timer timeit(config.logger, "setupLSPQueryBySymbolInFiles");
+    ENFORCE(symbol.exists());
+    return typechecker.query(core::lsp::Query::createSymbolQuery(symbol), frefs);
+}
+
+LSPQueryResult LSPQuery::bySymbol(const LSPConfiguration &config, LSPTypecheckerInterface &typechecker,
+                                  core::SymbolRef symbol) {
+    Timer timeit(config.logger, "setupLSPQueryBySymbol");
+    ENFORCE(symbol.exists());
+    vector<core::FileRef> frefs;
+    const core::GlobalState &gs = typechecker.state();
+    const core::NameHash symNameHash(gs, symbol.name(gs));
+    // Locate files that contain the same Name as the symbol. Is an overapproximation, but a good first filter.
+    int i = -1;
+    for (auto &file : typechecker.state().getFiles()) {
+        i++;
+        if (file == nullptr) {
+            continue;
+        }
+
+        ENFORCE(file->getFileHash() != nullptr);
+        const auto &hash = *file->getFileHash();
+        const auto &usedSends = hash.usages.sends;
+        const auto &usedSymbolNames = hash.usages.symbols;
+        auto ref = core::FileRef(i);
+
+        const bool fileIsValid = ref.exists() && ref.data(gs).sourceType == core::File::Type::Normal;
+        if (fileIsValid &&
+            (std::find(usedSends.begin(), usedSends.end(), symNameHash) != usedSends.end() ||
+             std::find(usedSymbolNames.begin(), usedSymbolNames.end(), symNameHash) != usedSymbolNames.end())) {
+            frefs.emplace_back(ref);
+        }
+    }
+
+    return typechecker.query(core::lsp::Query::createSymbolQuery(symbol), frefs);
+}
+
+} // namespace sorbet::realmain::lsp

main/lsp/LSPQuery.h

@@ -0,0 +1,25 @@
+#ifndef SORBET_LSP_QUERY_H
+#define SORBET_LSP_QUERY_H
+
+#include "main/lsp/LSPTypechecker.h"
+
+namespace sorbet::realmain::lsp {
+
+class LSPQuery {
+public:
+    static std::vector<std::unique_ptr<core::lsp::QueryResponse>>
+    filterAndDedup(const core::GlobalState &gs,
+                   const std::vector<std::unique_ptr<core::lsp::QueryResponse>> &queryResponses);
+
+    static LSPQueryResult byLoc(const LSPConfiguration &config, LSPTypecheckerInterface &typechecker,
+                                std::string_view uri, const Position &pos, LSPMethod forMethod,
+                                bool errorIfFileIsUntyped = true);
+    static LSPQueryResult bySymbolInFiles(const LSPConfiguration &config, LSPTypecheckerInterface &typechecker,
+                                          core::SymbolRef symbol, std::vector<core::FileRef> frefs);
+    static LSPQueryResult bySymbol(const LSPConfiguration &config, LSPTypecheckerInterface &typechecker,
+                                   core::SymbolRef symbol);
+};
+
+} // namespace sorbet::realmain::lsp
+
+#endif