Add find references support for methods

lib/ruby_indexer/lib/ruby_indexer/reference_finder.rb

@@ -5,6 +5,37 @@ module RubyIndexer
   class ReferenceFinder
     extend T::Sig
 
+    class Target
+      # def initialize
+      # end
+    end
+
+    class ConstTarget < Target
+      extend T::Sig
+
+      sig { params(fully_qualified_name: String).void }
+      def initialize(fully_qualified_name)
+        super()
+        @fully_qualified_name = fully_qualified_name
+      end
+
+      sig { returns(String) }
+      attr_reader :fully_qualified_name
+    end
+
+    class MethodTarget < Target
+      extend T::Sig
+
+      sig { params(method_name: String).void }
+      def initialize(method_name)
+        super()
+        @method_name = method_name
+      end
+
+      sig { returns(String) }
+      attr_reader :method_name
+    end
+
     class Reference
       extend T::Sig
 
@@ -27,15 +58,15 @@ def initialize(name, location, declaration:)
 
     sig do
       params(
-        fully_qualified_name: String,
+        target: Target,
         index: RubyIndexer::Index,
         dispatcher: Prism::Dispatcher,
         parse_result: Prism::ParseResult,
         include_declarations: T::Boolean,
       ).void
     end
-    def initialize(fully_qualified_name, index, dispatcher, parse_result, include_declarations: true)
-      @fully_qualified_name = fully_qualified_name
+    def initialize(target, index, dispatcher, parse_result, include_declarations: true)
+      @target = target
       @index = index
       @include_declarations = include_declarations
       @inside_def = T.let(false, T::Boolean)
@@ -64,6 +95,7 @@ def initialize(fully_qualified_name, index, dispatcher, parse_result, include_de
         :on_constant_or_write_node_enter,
         :on_constant_and_write_node_enter,
         :on_constant_operator_write_node_enter,
+        :on_call_node_enter,
       )
     end
 
@@ -80,7 +112,7 @@ def on_class_node_enter(node)
       name = constant_path.slice
       nesting = actual_nesting(name)
 
-      if nesting.join("::") == @fully_qualified_name
+      if @target.is_a?(ConstTarget) && nesting.join("::") == @target.fully_qualified_name
         @references << Reference.new(name, constant_path.location, declaration: true)
       end
 
@@ -98,7 +130,7 @@ def on_module_node_enter(node)
       name = constant_path.slice
       nesting = actual_nesting(name)
 
-      if nesting.join("::") == @fully_qualified_name
+      if @target.is_a?(ConstTarget) && nesting.join("::") == @target.fully_qualified_name
         @references << Reference.new(name, constant_path.location, declaration: true)
       end
 
@@ -215,6 +247,10 @@ def on_constant_operator_write_node_enter(node)
 
     sig { params(node: Prism::DefNode).void }
     def on_def_node_enter(node)
+      if @target.is_a?(MethodTarget) && (name = node.name.to_s) == @target.method_name
+        @references << Reference.new(name, node.name_loc, declaration: true)
+      end
+
       if node.receiver.is_a?(Prism::SelfNode)
         @stack << "<Class:#{@stack.last}>"
       end
@@ -227,6 +263,13 @@ def on_def_node_leave(node)
       end
     end
 
+    sig { params(node: Prism::CallNode).void }
+    def on_call_node_enter(node)
+      if @target.is_a?(MethodTarget) && (name = node.name.to_s) == @target.method_name
+        @references << Reference.new(name, T.must(node.message_loc), declaration: false)
+      end
+    end
+
     private
 
     sig { params(name: String).returns(T::Array[String]) }
@@ -245,13 +288,15 @@ def actual_nesting(name)
 
     sig { params(name: String, location: Prism::Location).void }
     def collect_constant_references(name, location)
+      return unless @target.is_a?(ConstTarget)
+
       entries = @index.resolve(name, @stack)
       return unless entries
 
       previous_reference = @references.last
 
       entries.each do |entry|
-        next unless entry.name == @fully_qualified_name
+        next unless entry.name == @target.fully_qualified_name
 
         # When processing a class/module declaration, we eagerly handle the constant reference. To avoid duplicates,
         # when we find the constant node defining the namespace, then we have to check if it wasn't already added

lib/ruby_indexer/test/reference_finder_test.rb

@@ -6,7 +6,8 @@
 module RubyIndexer
   class ReferenceFinderTest < Minitest::Test
     def test_finds_constant_references
-      refs = find_references("Foo::Bar", <<~RUBY)
+      target = ReferenceFinder::ConstTarget.new("Foo::Bar")
+      refs = find_references(target, <<~RUBY)
         module Foo
           class Bar
           end
@@ -27,15 +28,61 @@ class Bar
       assert_equal(8, refs[2].location.start_line)
     end
 
+    def test_finds_method_references
+      target = ReferenceFinder::MethodTarget.new("foo")
+      refs = find_references(target, <<~RUBY)
+        class Bar
+          def foo
+          end
+
+          def baz
+            foo
+          end
+        end
+      RUBY
+
+      assert_equal("foo", refs[0].name)
+      assert_equal(2, refs[0].location.start_line)
+
+      assert_equal("foo", refs[1].name)
+      assert_equal(6, refs[1].location.start_line)
+    end
+
+    def test_does_not_mismatch_on_attrs_readers_and_writers
+      target = ReferenceFinder::MethodTarget.new("foo")
+      refs = find_references(target, <<~RUBY)
+        class Bar
+          def foo
+          end
+
+          def foo=(value)
+          end
+
+          def baz
+            self.foo = 1
+            self.foo
+          end
+        end
+      RUBY
+
+      assert_equal(2, refs.size)
+
+      assert_equal("foo", refs[0].name)
+      assert_equal(2, refs[0].location.start_line)
+
+      assert_equal("foo", refs[1].name)
+      assert_equal(10, refs[1].location.start_line)
+    end
+
     private
 
-    def find_references(fully_qualified_name, source)
+    def find_references(target, source)
       file_path = "/fake.rb"
       index = Index.new
       index.index_single(IndexablePath.new(nil, file_path), source)
       parse_result = Prism.parse(source)
       dispatcher = Prism::Dispatcher.new
-      finder = ReferenceFinder.new(fully_qualified_name, index, dispatcher, parse_result)
+      finder = ReferenceFinder.new(target, index, dispatcher, parse_result)
       dispatcher.visit(parse_result.value)
       finder.references
     end

lib/ruby_lsp/requests/references.rb

@@ -35,7 +35,13 @@ def perform
         node_context = RubyDocument.locate(
           @document.parse_result.value,
           char_position,
-          node_types: [Prism::ConstantReadNode, Prism::ConstantPathNode, Prism::ConstantPathTargetNode],
+          node_types: [
+            Prism::ConstantReadNode,
+            Prism::ConstantPathNode,
+            Prism::ConstantPathTargetNode,
+            Prism::CallNode,
+            Prism::DefNode,
+          ],
         )
         target = node_context.node
         parent = node_context.parent
@@ -51,16 +57,17 @@ def perform
 
         target = T.cast(
           target,
-          T.any(Prism::ConstantReadNode, Prism::ConstantPathNode, Prism::ConstantPathTargetNode),
+          T.any(
+            Prism::ConstantReadNode,
+            Prism::ConstantPathNode,
+            Prism::ConstantPathTargetNode,
+            Prism::CallNode,
+            Prism::DefNode,
+          ),
         )
 
-        name = constant_name(target)
-        return @locations unless name
-
-        entries = @global_state.index.resolve(name, node_context.nesting)
-        return @locations unless entries
-
-        fully_qualified_name = T.must(entries.first).name
+        reference_target = create_reference_target(target, node_context)
+        return @locations unless reference_target
 
         Dir.glob(File.join(@global_state.workspace_path, "**/*.rb")).each do |path|
           uri = URI::Generic.from_path(path: path)
@@ -69,11 +76,11 @@ def perform
           next if @store.key?(uri)
 
           parse_result = Prism.parse_file(path)
-          collect_references(fully_qualified_name, parse_result, uri)
+          collect_references(reference_target, parse_result, uri)
         end
 
         @store.each do |_uri, document|
-          collect_references(fully_qualified_name, document.parse_result, document.uri)
+          collect_references(reference_target, document.parse_result, document.uri)
         end
 
         @locations
@@ -83,15 +90,43 @@ def perform
 
       sig do
         params(
-          fully_qualified_name: String,
+          target: T.any(
+            Prism::ConstantReadNode,
+            Prism::ConstantPathNode,
+            Prism::ConstantPathTargetNode,
+            Prism::CallNode,
+            Prism::DefNode,
+          ),
+          node_context: NodeContext,
+        ).returns(T.nilable(RubyIndexer::ReferenceFinder::Target))
+      end
+      def create_reference_target(target, node_context)
+        case target
+        when Prism::ConstantReadNode, Prism::ConstantPathNode, Prism::ConstantPathTargetNode
+          name = constant_name(target)
+          return unless name
+
+          entries = @global_state.index.resolve(name, node_context.nesting)
+          return unless entries
+
+          fully_qualified_name = T.must(entries.first).name
+          RubyIndexer::ReferenceFinder::ConstTarget.new(fully_qualified_name)
+        when Prism::CallNode, Prism::DefNode
+          RubyIndexer::ReferenceFinder::MethodTarget.new(target.name.to_s)
+        end
+      end
+
+      sig do
+        params(
+          target: RubyIndexer::ReferenceFinder::Target,
           parse_result: Prism::ParseResult,
           uri: URI::Generic,
         ).void
       end
-      def collect_references(fully_qualified_name, parse_result, uri)
+      def collect_references(target, parse_result, uri)
         dispatcher = Prism::Dispatcher.new
         finder = RubyIndexer::ReferenceFinder.new(
-          fully_qualified_name,
+          target,
           @global_state.index,
           dispatcher,
           parse_result,

lib/ruby_lsp/requests/rename.rb

@@ -60,7 +60,8 @@ def perform
         return unless entries
 
         fully_qualified_name = T.must(entries.first).name
-        changes = collect_text_edits(fully_qualified_name, name)
+        reference_target = RubyIndexer::ReferenceFinder::ConstTarget.new(fully_qualified_name)
+        changes = collect_text_edits(reference_target, name)
 
         # If the client doesn't support resource operations, such as renaming files, then we can only return the basic
         # text changes
@@ -124,8 +125,13 @@ def collect_file_renames(fully_qualified_name, document_changes)
         end
       end
 
-      sig { params(fully_qualified_name: String, name: String).returns(T::Hash[String, T::Array[Interface::TextEdit]]) }
-      def collect_text_edits(fully_qualified_name, name)
+      sig do
+        params(
+          target: RubyIndexer::ReferenceFinder::Target,
+          name: String,
+        ).returns(T::Hash[String, T::Array[Interface::TextEdit]])
+      end
+      def collect_text_edits(target, name)
         changes = {}
 
         Dir.glob(File.join(@global_state.workspace_path, "**/*.rb")).each do |path|
@@ -135,12 +141,12 @@ def collect_text_edits(fully_qualified_name, name)
           next if @store.key?(uri)
 
           parse_result = Prism.parse_file(path)
-          edits = collect_changes(fully_qualified_name, parse_result, name, uri)
+          edits = collect_changes(target, parse_result, name, uri)
           changes[uri.to_s] = edits unless edits.empty?
         end
 
         @store.each do |uri, document|
-          edits = collect_changes(fully_qualified_name, document.parse_result, name, document.uri)
+          edits = collect_changes(target, document.parse_result, name, document.uri)
           changes[uri] = edits unless edits.empty?
         end
 
@@ -149,16 +155,16 @@ def collect_text_edits(fully_qualified_name, name)
 
       sig do
         params(
-          fully_qualified_name: String,
+          target: RubyIndexer::ReferenceFinder::Target,
           parse_result: Prism::ParseResult,
           name: String,
           uri: URI::Generic,
         ).returns(T::Array[Interface::TextEdit])
       end
-      def collect_changes(fully_qualified_name, parse_result, name, uri)
+      def collect_changes(target, parse_result, name, uri)
         dispatcher = Prism::Dispatcher.new
         finder = RubyIndexer::ReferenceFinder.new(
-          fully_qualified_name,
+          target,
           @global_state.index,
           dispatcher,
           parse_result,