Add method support to References request (#2650)

* Add find references support for methods

* Format on one line

* Delete commented-out code

* Explain matching

* Fix tests

* Count matches in test

* Add test for matching writers

* PR feedback

* Mark Target as abstract

* More tests

Author: Andy Waite

lib/ruby_indexer/lib/ruby_indexer/reference_finder.rb

@@ -5,6 +5,38 @@ module RubyIndexer
   class ReferenceFinder
     extend T::Sig
 
+    class Target
+      extend T::Helpers
+
+      abstract!
+    end
+
+    class ConstTarget < Target
+      extend T::Sig
+
+      sig { returns(String) }
+      attr_reader :fully_qualified_name
+
+      sig { params(fully_qualified_name: String).void }
+      def initialize(fully_qualified_name)
+        super()
+        @fully_qualified_name = fully_qualified_name
+      end
+    end
+
+    class MethodTarget < Target
+      extend T::Sig
+
+      sig { returns(String) }
+      attr_reader :method_name
+
+      sig { params(method_name: String).void }
+      def initialize(method_name)
+        super()
+        @method_name = method_name
+      end
+    end
+
     class Reference
       extend T::Sig
 
@@ -27,14 +59,14 @@ def initialize(name, location, declaration:)
 
     sig do
       params(
-        fully_qualified_name: String,
+        target: Target,
         index: RubyIndexer::Index,
         dispatcher: Prism::Dispatcher,
         include_declarations: T::Boolean,
       ).void
     end
-    def initialize(fully_qualified_name, index, dispatcher, include_declarations: true)
-      @fully_qualified_name = fully_qualified_name
+    def initialize(target, index, dispatcher, include_declarations: true)
+      @target = target
       @index = index
       @include_declarations = include_declarations
       @stack = T.let([], T::Array[String])
@@ -62,6 +94,7 @@ def initialize(fully_qualified_name, index, dispatcher, include_declarations: tr
         :on_constant_or_write_node_enter,
         :on_constant_and_write_node_enter,
         :on_constant_operator_write_node_enter,
+        :on_call_node_enter,
       )
     end
 
@@ -78,7 +111,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
 
@@ -96,7 +129,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
 
@@ -213,6 +246,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
@@ -225,6 +262,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]) }
@@ -243,13 +287,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,7 @@
 module RubyIndexer
   class ReferenceFinderTest < Minitest::Test
     def test_finds_constant_references
-      refs = find_references("Foo::Bar", <<~RUBY)
+      refs = find_const_references("Foo::Bar", <<~RUBY)
         module Foo
           class Bar
           end
@@ -28,7 +28,7 @@ class Bar
     end
 
     def test_finds_constant_references_inside_singleton_contexts
-      refs = find_references("Foo::<Class:Foo>::Bar", <<~RUBY)
+      refs = find_const_references("Foo::<Class:Foo>::Bar", <<~RUBY)
         class Foo
           class << self
             class Bar
@@ -47,7 +47,7 @@ class Bar
     end
 
     def test_finds_top_level_constant_references
-      refs = find_references("Bar", <<~RUBY)
+      refs = find_const_references("Bar", <<~RUBY)
         class Bar
         end
 
@@ -70,15 +70,171 @@ class << self
       assert_equal(8, refs[2].location.start_line)
     end
 
+    def test_finds_method_references
+      refs = find_method_references("foo", <<~RUBY)
+        class Bar
+          def foo
+          end
+
+          def baz
+            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(6, refs[1].location.start_line)
+    end
+
+    def test_does_not_mismatch_on_readers_and_writers
+      refs = find_method_references("foo", <<~RUBY)
+        class Bar
+          def foo
+          end
+
+          def foo=(value)
+          end
+
+          def baz
+            self.foo = 1
+            self.foo
+          end
+        end
+      RUBY
+
+      # We want to match `foo` but not `foo=`
+      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
+
+    def test_matches_writers
+      refs = find_method_references("foo=", <<~RUBY)
+        class Bar
+          def foo
+          end
+
+          def foo=(value)
+          end
+
+          def baz
+            self.foo = 1
+            self.foo
+          end
+        end
+      RUBY
+
+      # We want to match `foo=` but not `foo`
+      assert_equal(2, refs.size)
+
+      assert_equal("foo=", refs[0].name)
+      assert_equal(5, refs[0].location.start_line)
+
+      assert_equal("foo=", refs[1].name)
+      assert_equal(9, refs[1].location.start_line)
+    end
+
+    def test_find_inherited_methods
+      refs = find_method_references("foo", <<~RUBY)
+        class Bar
+          def foo
+          end
+        end
+
+        class Baz < Bar
+          super.foo
+        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(7, refs[1].location.start_line)
+    end
+
+    def test_finds_methods_created_in_mixins
+      refs = find_method_references("foo", <<~RUBY)
+        module Mixin
+          def foo
+          end
+        end
+
+        class Bar
+          include Mixin
+        end
+
+        Bar.foo
+      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
+
+    def test_finds_singleton_methods
+      # The current implementation matches on both `Bar.foo` and `Bar#foo` even though they are different
+
+      refs = find_method_references("foo", <<~RUBY)
+        class Bar
+          class << self
+            def foo
+            end
+          end
+
+          def foo
+          end
+        end
+
+        Bar.foo
+      RUBY
+
+      assert_equal(3, refs.size)
+
+      assert_equal("foo", refs[0].name)
+      assert_equal(3, refs[0].location.start_line)
+
+      assert_equal("foo", refs[1].name)
+      assert_equal(7, refs[1].location.start_line)
+
+      assert_equal("foo", refs[2].name)
+      assert_equal(11, refs[2].location.start_line)
+    end
+
     private
 
-    def find_references(fully_qualified_name, source)
+    def find_const_references(const_name, source)
+      target = ReferenceFinder::ConstTarget.new(const_name)
+      find_references(target, source)
+    end
+
+    def find_method_references(method_name, source)
+      target = ReferenceFinder::MethodTarget.new(method_name)
+      find_references(target, source)
+    end
+
+    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)
+      finder = ReferenceFinder.new(target, index, dispatcher)
       dispatcher.visit(parse_result.value)
       finder.references
     end