[LSP] Go to implementation (sorbet#4598)

* Implement go to implementation for abstract classes

* New use cases for the go toimplentation:

- Go to implementation from abstract class site

- Go to implementation from abstract class reference

- Go to implementation from abstract method call site

Also change updates tests

* Replace has_value() with value_or(false) for "Go to implementation" experiment feature in lsp test

* Move getSubclasses to the ClassOrModuleRef

* Address PR feedback

* Run linter

* Remove allLocs

* Fix remaining use cases

* Run linter

* Fix server capabilities

* Fix test

* Another attempt to fix recorded tests

* Add additional safety check

* Add tests for the "Go to Implementations"

* Run formatting

* Address comments

* PR comments: part 2

* Move local functions into anonymous namespace

.gitignore

@@ -9,6 +9,7 @@
 
 .idea/
 .vscode/
+.vim/
 
 /website/lib/core/metadata.js
 /website/lib/core/MetadataBlog.js

README.md

@@ -659,6 +659,33 @@ class Project::Foo
 end
 ```
 
+#### Testing "Go to Implementation"
+
+Testing the "Go to Implementation" feature is really similar to the testing techniques of the "Go to Type Definition".
+
+```ruby
+module A
+#      ^ find-implementation: A
+  extend T::Sig
+  extend T::Helpers
+  interface!
+end
+
+ class B
+#^^^^^^^ implementation: A
+  extend T::Sig
+  include A
+#         ^ find-implementation: A
+end
+```
+
+There are two types of assertions:
+
+1. `find-implementation: <symbol>` means make a "Go to Implementation" request here. `<symbol>` marks the symbol name we are looking for.
+2. `implementation: <symbol>` marks the location which should be returned for the "Go to Implementation" call for a given `<symbol>`
+
+If the request returns multiple locations, you should mark all of them with `implementation: <symbol>`
+
 #### Testing rename constant
 
 To write a test for renaming constants, you need to make at least two files:

core/SymbolRef.h

@@ -74,6 +74,7 @@ class ClassOrModuleRef final {
     SymbolData dataAllowingNone(GlobalState &gs) const;
     ConstSymbolData data(const GlobalState &gs) const;
     ConstSymbolData dataAllowingNone(const GlobalState &gs) const;
+    std::vector<core::ClassOrModuleRef> getSubclasses(const core::GlobalState &gs, bool withSelf = true);
 
     bool operator==(const ClassOrModuleRef &rhs) const;
     bool operator!=(const ClassOrModuleRef &rhs) const;

core/Symbols.cc

@@ -12,14 +12,41 @@
 #include "core/errors/internal.h"
 #include "core/hashing/hashing.h"
 #include <string>
+#include <vector>
 
 template class std::vector<sorbet::core::TypeAndOrigins>;
 template class std::vector<std::pair<sorbet::core::NameRef, sorbet::core::SymbolRef>>;
 template class std::vector<const sorbet::core::Symbol *>;
 
 namespace sorbet::core {
+
 using namespace std;
 
+namespace {
+
+// Checks if s is a subclass of root or contains root as a mixin, and updates visited and memoized vectors.
+bool isSubclassOrMixin(const core::GlobalState &gs, core::ClassOrModuleRef root, core::ClassOrModuleRef s,
+                       std::vector<bool> &memoized, std::vector<bool> &visited) {
+    // don't visit the same class twice
+    if (visited[s.id()] == true) {
+        return memoized[s.id()];
+    }
+    visited[s.id()] = true;
+
+    for (auto a : s.data(gs)->mixins()) {
+        if (a == root) {
+            memoized[s.id()] = true;
+            return true;
+        }
+    }
+    if (s.data(gs)->superClass().exists()) {
+        memoized[s.id()] = isSubclassOrMixin(gs, root, s.data(gs)->superClass(), memoized, visited);
+    }
+
+    return memoized[s.id()];
+}
+} // namespace
+
 namespace {
 constexpr string_view COLON_SEPARATOR = "::"sv;
 constexpr string_view HASH_SEPARATOR = "#"sv;
@@ -49,6 +76,26 @@ bool ClassOrModuleRef::operator!=(const ClassOrModuleRef &rhs) const {
     return rhs._id != this->_id;
 }
 
+// Returns all subclasses of ClassOrModuleRef (including itself)
+vector<core::ClassOrModuleRef> ClassOrModuleRef::getSubclasses(const core::GlobalState &gs, bool withSelf) {
+    vector<bool> memoized(gs.classAndModulesUsed());
+    vector<bool> visited(gs.classAndModulesUsed());
+    memoized[this->id()] = true;
+    visited[this->id()] = true;
+
+    vector<core::ClassOrModuleRef> subclasses;
+    for (u4 i = 1; i < gs.classAndModulesUsed(); ++i) {
+        auto s = core::ClassOrModuleRef(gs, i);
+        if (!withSelf && s == *this) {
+            continue;
+        }
+        if (isSubclassOrMixin(gs, *this, s, memoized, visited)) {
+            subclasses.emplace_back(s);
+        }
+    }
+    return subclasses;
+}
+
 bool MethodRef::operator==(const MethodRef &rhs) const {
     return rhs._id == this->_id;
 }

main/lsp/LSPPreprocessor.cc

@@ -198,6 +198,9 @@ unique_ptr<LSPTask> LSPPreprocessor::getTaskForMessage(LSPMessage &msg) {
                                                       move(get<unique_ptr<TextDocumentPositionParams>>(rawParams)));
             case LSPMethod::TextDocumentReferences:
                 return make_unique<ReferencesTask>(*config, id, move(get<unique_ptr<ReferenceParams>>(rawParams)));
+            case LSPMethod::TextDocumentImplementation:
+                return make_unique<GoToImplementationTask>(*config, id,
+                                                           move(get<unique_ptr<ImplementationParams>>(rawParams)));
             case LSPMethod::SorbetReadFile:
                 return make_unique<SorbetReadFileTask>(*config, id,
                                                        move(get<unique_ptr<TextDocumentIdentifier>>(rawParams)));

main/lsp/LSPTask.cc

@@ -83,6 +83,8 @@ ConstExprStr LSPTask::methodString() const {
             return "textDocument.typeDefinition";
         case LSPMethod::WorkspaceSymbol:
             return "workspace.symbol";
+        case LSPMethod::TextDocumentImplementation:
+            return "textDocument.implementation";
     }
 }
 

main/lsp/requests/go_to_implementation.cc

@@ -0,0 +1,158 @@
+#include "main/lsp/requests/go_to_implementation.h"
+#include "core/lsp/QueryResponse.h"
+#include "main/lsp/json_types.h"
+#include "main/lsp/lsp.h"
+#include <cstddef>
+#include <memory>
+
+using namespace std;
+
+namespace sorbet::realmain::lsp {
+
+namespace {
+struct MethodImplementationResults {
+    vector<core::Loc> locations;
+    unique_ptr<ResponseError> error;
+};
+
+unique_ptr<ResponseError> makeInvalidParamsError(std::string error) {
+    return make_unique<ResponseError>((int)LSPErrorCodes::InvalidParams, error);
+}
+
+const MethodImplementationResults findMethodImplementations(const core::GlobalState &gs, core::SymbolRef method) {
+    MethodImplementationResults res;
+    if (!method.data(gs)->isMethod() || !method.data(gs)->isAbstract()) {
+        res.error = makeInvalidParamsError(
+            "Go to implementation can be used only for methods or references of abstract classes");
+        return res;
+    }
+
+    vector<core::Loc> locations;
+    auto owner = method.data(gs)->owner;
+    if (!owner.isClassOrModule()) {
+        res.error = makeInvalidParamsError("Abstract method can only be inside a class or module");
+        return res;
+    }
+
+    auto owningClassSymbolRef = owner.asClassOrModuleRef();
+    auto childClasses = owningClassSymbolRef.getSubclasses(gs, false);
+    auto methodName = method.data(gs)->name;
+    for (const auto &childClass : childClasses) {
+        auto methodImplementation = childClass.data(gs)->findMember(gs, methodName);
+        locations.push_back(methodImplementation.data(gs)->loc());
+    }
+
+    res.locations = locations;
+    return res;
+}
+
+core::SymbolRef findOverridedMethod(const core::GlobalState &gs, const core::SymbolRef method) {
+    auto ownerClass = method.data(gs)->owner.asClassOrModuleRef();
+
+    for (auto mixin : ownerClass.data(gs)->mixins()) {
+        if (!mixin.data(gs)->isClassOrModule() && !mixin.data(gs)->isAbstract()) {
+            continue;
+        }
+        return mixin.data(gs)->findMember(gs, method.data(gs)->name);
+    }
+    return core::Symbols::noSymbol();
+}
+} // namespace
+
+GoToImplementationTask::GoToImplementationTask(const LSPConfiguration &config, MessageId id,
+                                               std::unique_ptr<ImplementationParams> params)
+    : LSPRequestTask(config, move(id), LSPMethod::TextDocumentImplementation), params(move(params)) {}
+
+unique_ptr<ResponseMessage> GoToImplementationTask::runRequest(LSPTypecheckerDelegate &typechecker) {
+    auto response = make_unique<ResponseMessage>("2.0", id, LSPMethod::TextDocumentImplementation);
+
+    const core::GlobalState &gs = typechecker.state();
+    auto queryResult =
+        queryByLoc(typechecker, params->textDocument->uri, *params->position, LSPMethod::TextDocumentImplementation);
+
+    if (queryResult.error) {
+        // An error happened while setting up the query.
+        response->error = move(queryResult.error);
+        return response;
+    }
+
+    if (queryResult.responses.empty()) {
+        return response;
+    }
+
+    vector<unique_ptr<Location>> result;
+    auto queryResponse = move(queryResult.responses[0]);
+    if (auto def = queryResponse->isDefinition()) {
+        // User called "Go to Implementation" from the abstract function definition
+        core::SymbolRef method = def->symbol;
+        if (!method.data(gs)->isMethod()) {
+            response->error = make_unique<ResponseError>(
+                (int)LSPErrorCodes::InvalidParams,
+                "Go to implementation can be used only for methods or references of abstract classes");
+            return response;
+        }
+
+        core::SymbolRef overridedMethod = method;
+        if (method.data(gs)->isOverride()) {
+            overridedMethod = findOverridedMethod(gs, method);
+        }
+        auto locationsOrError = findMethodImplementations(gs, overridedMethod);
+
+        if (locationsOrError.error != nullptr) {
+            response->error = move(locationsOrError.error);
+            return response;
+        } else {
+            for (const auto &location : locationsOrError.locations) {
+                addLocIfExists(gs, result, location);
+            }
+        }
+    } else if (auto constant = queryResponse->isConstant()) {
+        // User called "Go to Implementation" from the abstract class reference
+        auto classSymbol = constant->symbol;
+
+        if (!classSymbol.data(gs)->isClassOrModule() || !classSymbol.data(gs)->isClassOrModuleAbstract()) {
+            response->error = make_unique<ResponseError>(
+                (int)LSPErrorCodes::InvalidParams,
+                "Go to implementation can be used only for methods or references of abstract classes");
+            return response;
+        }
+
+        auto classOrModuleRef = classSymbol.asClassOrModuleRef();
+        auto childClasses = classOrModuleRef.getSubclasses(gs, false);
+        for (const auto &childClass : childClasses) {
+            for (auto loc : childClass.data(gs)->locs()) {
+                addLocIfExists(gs, result, loc);
+            }
+        }
+
+    } else if (auto send = queryResponse->isSend()) {
+        auto mainResponse = move(send->dispatchResult->main);
+
+        // User called "Go to Implementation" from the abstract function call
+        if (mainResponse.errors.size() != 0) {
+            response->error = makeInvalidParamsError("Failed to fetch implementations");
+            return response;
+        }
+
+        auto calledMethod = mainResponse.method;
+        core::SymbolRef overridedMethod = calledMethod;
+        if (calledMethod.data(gs)->isOverride()) {
+            overridedMethod = findOverridedMethod(gs, overridedMethod);
+        }
+
+        auto locationsOrError = findMethodImplementations(gs, overridedMethod);
+
+        if (locationsOrError.error != nullptr) {
+            response->error = move(locationsOrError.error);
+            return response;
+        } else {
+            for (const auto &location : locationsOrError.locations) {
+                addLocIfExists(gs, result, location);
+            }
+        }
+    }
+
+    response->result = move(result);
+    return response;
+}
+} // namespace sorbet::realmain::lsp

main/lsp/requests/go_to_implementation.h

@@ -0,0 +1,20 @@
+#ifndef RUBY_TYPER_LSP_REQUESTS_GO_TO_IMPLEMENTATION_H
+#define RUBY_TYPER_LSP_REQUESTS_GO_TO_IMPLEMENTATION_H
+
+#include "main/lsp/LSPTask.h"
+
+namespace sorbet::realmain::lsp {
+class ImplementationParams;
+
+class GoToImplementationTask final : public LSPRequestTask {
+    std::unique_ptr<ImplementationParams> params;
+
+public:
+    GoToImplementationTask(const LSPConfiguration &config, MessageId id, std::unique_ptr<ImplementationParams> params);
+
+    std::unique_ptr<ResponseMessage> runRequest(LSPTypecheckerDelegate &typechecker) override;
+};
+
+} // namespace sorbet::realmain::lsp
+
+#endif

main/lsp/requests/initialize.cc

@@ -27,6 +27,7 @@ unique_ptr<ResponseMessage> InitializeTask::runRequest(LSPTypecheckerDelegate &t
     serverCap->documentHighlightProvider = opts.lspDocumentHighlightEnabled;
     serverCap->hoverProvider = true;
     serverCap->referencesProvider = true;
+    serverCap->implementationProvider = opts.lspGoToImplementationEnabled;
     serverCap->documentFormattingProvider = rubyfmt_enabled && opts.lspDocumentFormatRubyfmtEnabled;
 
     auto codeActionProvider = make_unique<CodeActionOptions>();

main/lsp/requests/rename.cc

@@ -35,45 +35,6 @@ bool isValidRenameLocation(const core::SymbolRef &symbol, const core::GlobalStat
     return true;
 }
 
-// Checks if s is a subclass of root or contains root as a mixin, and updates visited and memoized vectors.
-bool isSubclassOrMixin(const core::GlobalState &gs, core::ClassOrModuleRef root, core::ClassOrModuleRef s,
-                       vector<bool> &memoized, vector<bool> &visited) {
-    // don't visit the same class twice
-    if (visited[s.id()] == true) {
-        return memoized[s.id()];
-    }
-    visited[s.id()] = true;
-
-    for (auto a : s.data(gs)->mixins()) {
-        if (a == root) {
-            memoized[s.id()] = true;
-            return true;
-        }
-    }
-    if (s.data(gs)->superClass().exists()) {
-        memoized[s.id()] = isSubclassOrMixin(gs, root, s.data(gs)->superClass(), memoized, visited);
-    }
-
-    return memoized[s.id()];
-}
-
-// Returns all subclasses of root (including root)
-vector<core::ClassOrModuleRef> getSubclasses(const core::GlobalState &gs, core::ClassOrModuleRef root) {
-    vector<bool> memoized(gs.classAndModulesUsed());
-    vector<bool> visited(gs.classAndModulesUsed());
-    memoized[root.id()] = true;
-    visited[root.id()] = true;
-
-    vector<core::ClassOrModuleRef> subclasses;
-    for (u4 i = 1; i < gs.classAndModulesUsed(); ++i) {
-        auto s = core::ClassOrModuleRef(gs, i);
-        if (isSubclassOrMixin(gs, root, s, memoized, visited)) {
-            subclasses.emplace_back(s);
-        }
-    }
-    return subclasses;
-}
-
 // Follow superClass links until we find the highest class that contains the given method. In other words we find the
 // "root" of the tree of classes that define a method.
 core::ClassOrModuleRef findRootClassWithMethod(const core::GlobalState &gs, core::ClassOrModuleRef klass,
@@ -131,7 +92,7 @@ void addSubclassRelatedMethods(const core::GlobalState &gs, core::MethodRef symb
     // method_class_hierarchy test case for an example).
     auto root = findRootClassWithMethod(gs, symbolClass, symbolData->name);
 
-    auto subclasses = getSubclasses(gs, root);
+    auto subclasses = root.getSubclasses(gs);
 
     // find the target method definition in each subclass
     for (auto c : subclasses) {
@@ -461,4 +422,4 @@ unique_ptr<ResponseMessage> RenameTask::runRequest(LSPTypecheckerDelegate &typec
     return response;
 }
 
-} // namespace sorbet::realmain::lsp
+} // namespace sorbet::realmain::lsp

main/lsp/requests/requests.h

@@ -8,6 +8,7 @@
 #include "main/lsp/requests/document_highlight.h"
 #include "main/lsp/requests/document_symbol.h"
 #include "main/lsp/requests/get_counters.h"
+#include "main/lsp/requests/go_to_implementation.h"
 #include "main/lsp/requests/hover.h"
 #include "main/lsp/requests/initialize.h"
 #include "main/lsp/requests/prepare_rename.h"