Merge pull request #4015 from Shopify/03-03-migrate_go_to_definition_to_use_rubydex

Partially migrate go to definition to use Rubydex

Author: vinistock

lib/ruby_lsp/internal.rb

@@ -31,6 +31,9 @@
 require "shellwords"
 require "set"
 
+# Rubydex LSP additions
+require "ruby_lsp/rubydex/definition"
+
 require "ruby-lsp"
 require "ruby_lsp/base_server"
 require "ruby_indexer/ruby_indexer"

lib/ruby_lsp/listeners/definition.rb

@@ -12,7 +12,7 @@ class Definition
       def initialize(response_builder, global_state, language_id, uri, node_context, dispatcher, sorbet_level) # rubocop:disable Metrics/ParameterLists
         @response_builder = response_builder
         @global_state = global_state
-        @index = global_state.index #: RubyIndexer::Index
+        @graph = global_state.graph #: Rubydex::Graph
         @type_inferrer = global_state.type_inferrer #: TypeInferrer
         @language_id = language_id
         @uri = uri
@@ -109,15 +109,15 @@ def on_constant_path_node_enter(node)
         name = RubyIndexer::Index.constant_name(node)
         return if name.nil?
 
-        find_in_index(name)
+        handle_constant_definition(name)
       end
 
       #: (Prism::ConstantReadNode node) -> void
       def on_constant_read_node_enter(node)
         name = RubyIndexer::Index.constant_name(node)
         return if name.nil?
 
-        find_in_index(name)
+        handle_constant_definition(name)
       end
 
       #: (Prism::GlobalVariableAndWriteNode node) -> void
@@ -152,32 +152,32 @@ def on_global_variable_write_node_enter(node)
 
       #: (Prism::InstanceVariableReadNode node) -> void
       def on_instance_variable_read_node_enter(node)
-        handle_instance_variable_definition(node.name.to_s)
+        handle_variable_definition(node.name.to_s)
       end
 
       #: (Prism::InstanceVariableWriteNode node) -> void
       def on_instance_variable_write_node_enter(node)
-        handle_instance_variable_definition(node.name.to_s)
+        handle_variable_definition(node.name.to_s)
       end
 
       #: (Prism::InstanceVariableAndWriteNode node) -> void
       def on_instance_variable_and_write_node_enter(node)
-        handle_instance_variable_definition(node.name.to_s)
+        handle_variable_definition(node.name.to_s)
       end
 
       #: (Prism::InstanceVariableOperatorWriteNode node) -> void
       def on_instance_variable_operator_write_node_enter(node)
-        handle_instance_variable_definition(node.name.to_s)
+        handle_variable_definition(node.name.to_s)
       end
 
       #: (Prism::InstanceVariableOrWriteNode node) -> void
       def on_instance_variable_or_write_node_enter(node)
-        handle_instance_variable_definition(node.name.to_s)
+        handle_variable_definition(node.name.to_s)
       end
 
       #: (Prism::InstanceVariableTargetNode node) -> void
       def on_instance_variable_target_node_enter(node)
-        handle_instance_variable_definition(node.name.to_s)
+        handle_variable_definition(node.name.to_s)
       end
 
       #: (Prism::SuperNode node) -> void
@@ -192,32 +192,32 @@ def on_forwarding_super_node_enter(node)
 
       #: (Prism::ClassVariableAndWriteNode node) -> void
       def on_class_variable_and_write_node_enter(node)
-        handle_class_variable_definition(node.name.to_s)
+        handle_variable_definition(node.name.to_s)
       end
 
       #: (Prism::ClassVariableOperatorWriteNode node) -> void
       def on_class_variable_operator_write_node_enter(node)
-        handle_class_variable_definition(node.name.to_s)
+        handle_variable_definition(node.name.to_s)
       end
 
       #: (Prism::ClassVariableOrWriteNode node) -> void
       def on_class_variable_or_write_node_enter(node)
-        handle_class_variable_definition(node.name.to_s)
+        handle_variable_definition(node.name.to_s)
       end
 
       #: (Prism::ClassVariableTargetNode node) -> void
       def on_class_variable_target_node_enter(node)
-        handle_class_variable_definition(node.name.to_s)
+        handle_variable_definition(node.name.to_s)
       end
 
       #: (Prism::ClassVariableReadNode node) -> void
       def on_class_variable_read_node_enter(node)
-        handle_class_variable_definition(node.name.to_s)
+        handle_variable_definition(node.name.to_s)
       end
 
       #: (Prism::ClassVariableWriteNode node) -> void
       def on_class_variable_write_node_enter(node)
-        handle_class_variable_definition(node.name.to_s)
+        handle_variable_definition(node.name.to_s)
       end
 
       private
@@ -257,106 +257,74 @@ def handle_super_node_definition
 
       #: (String name) -> void
       def handle_global_variable_definition(name)
-        entries = @index[name]
+        declaration = @graph[name]
+        return unless declaration
 
-        return unless entries
-
-        entries.each do |entry|
-          location = entry.location
-
-          @response_builder << Interface::Location.new(
-            uri: entry.uri.to_s,
-            range: Interface::Range.new(
-              start: Interface::Position.new(line: location.start_line - 1, character: location.start_column),
-              end: Interface::Position.new(line: location.end_line - 1, character: location.end_column),
-            ),
-          )
-        end
-      end
-
-      #: (String name) -> void
-      def handle_class_variable_definition(name)
-        type = @type_inferrer.infer_receiver_type(@node_context)
-        return unless type
-
-        entries = @index.resolve_class_variable(name, type.name)
-        return unless entries
-
-        entries.each do |entry|
-          @response_builder << Interface::Location.new(
-            uri: entry.uri.to_s,
-            range: range_from_location(entry.location),
-          )
-        end
-      rescue RubyIndexer::Index::NonExistingNamespaceError
-        # If by any chance we haven't indexed the owner, then there's no way to find the right declaration
+        declaration.definitions.each { |definition| @response_builder << definition.to_lsp_selection_location }
       end
 
+      # Handle class or instance variables. We collect all definitions across the ancestors of the type
+      #
       #: (String name) -> void
-      def handle_instance_variable_definition(name)
-        # Sorbet enforces that all instance variables be declared on typed strict or higher, which means it will be able
-        # to provide all features for them
+      def handle_variable_definition(name)
+        # Sorbet enforces that all variables be declared on typed strict or higher, which means it will be able to
+        # provide all features for them
         return if @sorbet_level.strict?
 
         type = @type_inferrer.infer_receiver_type(@node_context)
         return unless type
 
-        entries = @index.resolve_instance_variable(name, type.name)
-        return unless entries
+        owner = @graph[type.name]
+        return unless owner.is_a?(Rubydex::Namespace)
 
-        entries.each do |entry|
-          location = entry.location
+        owner.ancestors.each do |ancestor|
+          member = ancestor.member(name)
+          next unless member
 
-          @response_builder << Interface::Location.new(
-            uri: entry.uri.to_s,
-            range: Interface::Range.new(
-              start: Interface::Position.new(line: location.start_line - 1, character: location.start_column),
-              end: Interface::Position.new(line: location.end_line - 1, character: location.end_column),
-            ),
-          )
+          member.definitions.each { |definition| @response_builder << definition.to_lsp_selection_location }
         end
-      rescue RubyIndexer::Index::NonExistingNamespaceError
-        # If by any chance we haven't indexed the owner, then there's no way to find the right declaration
       end
 
       #: (String message, TypeInferrer::Type? receiver_type, ?inherited_only: bool) -> void
       def handle_method_definition(message, receiver_type, inherited_only: false)
-        methods = if receiver_type
-          @index.resolve_method(message, receiver_type.name, inherited_only: inherited_only)
+        declaration = if receiver_type
+          owner = @graph[receiver_type.name]
+          owner.find_member("#{message}()", only_inherited: inherited_only) if owner.is_a?(Rubydex::Namespace)
         end
 
-        # If the method doesn't have a receiver, or the guessed receiver doesn't have any matched candidates,
-        # then we provide a few candidates to jump to
-        # But we don't want to provide too many candidates, as it can be overwhelming
-        if receiver_type.nil? || (receiver_type.is_a?(TypeInferrer::GuessedType) && methods.nil?)
-          methods = @index[message]&.take(MAX_NUMBER_OF_DEFINITION_CANDIDATES_WITHOUT_RECEIVER)
+        # If the method doesn't have a receiver, or the guessed receiver doesn't have any matched candidates, then we
+        # provide a few candidates to jump to. However, we don't want to provide too many candidates, as it can be
+        # overwhelming
+        if receiver_type.nil? || (receiver_type.is_a?(TypeInferrer::GuessedType) && declaration.nil?)
+          declaration = @graph.search("##{message}()").take(MAX_NUMBER_OF_DEFINITION_CANDIDATES_WITHOUT_RECEIVER)
         end
 
-        return unless methods
+        return unless declaration
 
-        methods.each do |target_method|
-          uri = target_method.uri
-          full_path = uri.full_path
-          next if @sorbet_level.true_or_higher? && (!full_path || not_in_dependencies?(full_path))
+        Array(declaration).each do |decl|
+          decl.definitions.each do |definition|
+            location = definition.location
+            uri = URI(location.uri)
+            full_path = uri.full_path
+            next if @sorbet_level.true_or_higher? && (!full_path || not_in_dependencies?(full_path))
 
-          @response_builder << Interface::LocationLink.new(
-            target_uri: uri.to_s,
-            target_range: range_from_location(target_method.location),
-            target_selection_range: range_from_location(target_method.name_location),
-          )
+            @response_builder << Interface::LocationLink.new(
+              target_uri: uri.to_s,
+              target_range: definition.to_lsp_selection_range,
+              target_selection_range: definition.to_lsp_name_range || definition.to_lsp_selection_range,
+            )
+          end
         end
       end
 
       #: (Prism::StringNode node, Symbol message) -> void
       def handle_require_definition(node, message)
         case message
         when :require
-          entry = @index.search_require_paths(node.content).find do |uri|
-            uri.require_path == node.content
-          end
+          document = @graph.resolve_require_path(node.content, $LOAD_PATH)
 
-          if entry
-            candidate = entry.full_path
+          if document
+            candidate = URI(document.uri).full_path
 
             if candidate
               @response_builder << Interface::Location.new(
@@ -392,35 +360,33 @@ def handle_autoload_definition(node)
         constant_name = argument.value
         return unless constant_name
 
-        find_in_index(constant_name)
+        handle_constant_definition(constant_name)
       end
 
       #: (String value) -> void
-      def find_in_index(value)
-        entries = @index.resolve(value, @node_context.nesting)
-        return unless entries
+      def handle_constant_definition(value)
+        declaration = @graph.resolve_constant(value, @node_context.nesting)
+        return unless declaration
 
+        # [RUBYDEX] TODO: temporarily commented out until we have the visibility API
+        #
         # We should only allow jumping to the definition of private constants if the constant is defined in the same
         # namespace as the reference
-        first_entry = entries.first #: as !nil
-        return if first_entry.private? && first_entry.name != "#{@node_context.fully_qualified_name}::#{value}"
+        #
+        # return if declaration.private? && declaration.name != "#{@node_context.fully_qualified_name}::#{value}"
 
-        entries.each do |entry|
+        declaration.definitions.each do |definition|
           # If the project has Sorbet, then we only want to handle go to definition for constants defined in gems, as an
           # additional behavior on top of jumping to RBIs. The only sigil where Sorbet cannot handle constants is typed
           # ignore
-          uri = entry.uri
+          uri = URI(definition.location.uri)
           full_path = uri.full_path
 
           if !@sorbet_level.ignore? && (!full_path || not_in_dependencies?(full_path))
             next
           end
 
-          @response_builder << Interface::LocationLink.new(
-            target_uri: uri.to_s,
-            target_range: range_from_location(entry.location),
-            target_selection_range: range_from_location(entry.name_location),
-          )
+          @response_builder << definition.to_lsp_location_link
         end
       end
     end

lib/ruby_lsp/rubydex/definition.rb

@@ -0,0 +1,65 @@
+# typed: strict
+# frozen_string_literal: true
+
+module Rubydex
+  class Definition
+    #: () -> RubyLsp::Interface::LocationLink
+    def to_lsp_location_link
+      selection_range = to_lsp_selection_range
+
+      RubyLsp::Interface::LocationLink.new(
+        target_uri: location.uri,
+        target_range: selection_range,
+        target_selection_range: to_lsp_name_range || selection_range,
+      )
+    end
+
+    #: () -> RubyLsp::Interface::Range
+    def to_lsp_selection_range
+      loc = location
+
+      RubyLsp::Interface::Range.new(
+        start: RubyLsp::Interface::Position.new(line: loc.start_line, character: loc.start_column),
+        end: RubyLsp::Interface::Position.new(line: loc.end_line, character: loc.end_column),
+      )
+    end
+
+    #: () -> RubyLsp::Interface::Location
+    def to_lsp_selection_location
+      location = self.location
+
+      RubyLsp::Interface::Location.new(
+        uri: location.uri,
+        range: RubyLsp::Interface::Range.new(
+          start: RubyLsp::Interface::Position.new(line: location.start_line, character: location.start_column),
+          end: RubyLsp::Interface::Position.new(line: location.end_line, character: location.end_column),
+        ),
+      )
+    end
+
+    #: () -> RubyLsp::Interface::Range?
+    def to_lsp_name_range
+      loc = name_location
+      return unless loc
+
+      RubyLsp::Interface::Range.new(
+        start: RubyLsp::Interface::Position.new(line: loc.start_line, character: loc.start_column),
+        end: RubyLsp::Interface::Position.new(line: loc.end_line, character: loc.end_column),
+      )
+    end
+
+    #: () -> RubyLsp::Interface::Location?
+    def to_lsp_name_location
+      location = name_location
+      return unless location
+
+      RubyLsp::Interface::Location.new(
+        uri: location.uri,
+        range: RubyLsp::Interface::Range.new(
+          start: RubyLsp::Interface::Position.new(line: location.start_line, character: location.start_column),
+          end: RubyLsp::Interface::Position.new(line: location.end_line, character: location.end_column),
+        ),
+      )
+    end
+  end
+end