Use method receivers to narrow reference search results (#4091)

Author: vinistock

lib/ruby_lsp/requests/references.rb

@@ -145,7 +145,7 @@ def resolve_method_references(message, node_context, include_declarations)
         end
         return if declarations.empty?
 
-        collect_references(method_references_for(message), declarations, include_declarations)
+        collect_references(method_references_for(message, declarations), declarations, include_declarations)
       end
 
       # Handles instance and class variable references. Resolves the receiver type from the node context to locate
@@ -187,14 +187,28 @@ def handle_def_node_references(target, node_context, include_declarations)
         declaration = owner.find_member("#{method_name}()")
         return unless declaration
 
-        collect_references(method_references_for(method_name), [declaration], include_declarations)
+        collect_references(method_references_for(method_name, [declaration]), [declaration], include_declarations)
       end
 
-      # Method references in Rubydex are not yet resolved to specific declarations, so we filter from the global
-      # method references by name
-      #: (String) -> Array[Rubydex::MethodReference]
-      def method_references_for(method_name)
-        @graph.method_references.select { |reference| reference.name == method_name }
+      #: (String, Array[Rubydex::Declaration]) -> Array[Rubydex::MethodReference]
+      def method_references_for(method_name, declarations)
+        target_owner_names = declarations.map do |d|
+          d.owner #: as Rubydex::Namespace
+            .name
+        end
+
+        @graph.method_references.select do |reference|
+          next false unless reference.name == method_name
+
+          receiver = reference.receiver
+          next true if receiver.nil? || target_owner_names.empty?
+
+          if receiver.is_a?(Rubydex::Namespace)
+            receiver.ancestors.any? { |ancestor| target_owner_names.include?(ancestor.name) }
+          else
+            target_owner_names.include?(receiver.name)
+          end
+        end
       end
 
       #: (Enumerable[Rubydex::Reference] references, Array[Rubydex::Declaration] declarations, bool include_declarations) -> void

test/requests/references_test.rb

@@ -71,43 +71,72 @@ def bar
     assert_includes(ref_lines, 5) # reference
   end
 
-  def test_does_not_match_simple_names_when_method_receiver_is_known
-    source = <<~RUBY
-      class Foo
-        def self.bar
-        end
+  AMBIGUOUS_BAR_SOURCE = <<~RUBY
+    class Foo
+      def self.bar
       end
+    end
 
-      class Qux
-        def bar
-        end
+    class Qux
+      def bar
+      end
+    end
+
+    class Other
+      def self.bar
       end
+    end
 
-      it = Qux.new
-      it.bar
+    it = Qux.new
+    it.bar
 
-      Foo.bar
-    RUBY
+    Foo.bar
+    Other.bar
+  RUBY
 
-    # Cursor on `bar` in `Foo.bar` — the receiver is the `Foo` singleton, so only the `Foo.bar`
-    # declaration is a valid target. We must not surface `Qux#bar` as a candidate declaration.
-    refs = find_references(source, { line: 13, character: 4 }, include_declarations: true)
+  def test_filters_method_references_when_call_site_receiver_is_a_known_constant
+    refs = find_references(AMBIGUOUS_BAR_SOURCE, { line: 18, character: 4 }, include_declarations: true)
     ref_lines = refs.map { |r| r.range.start.line }
+
     assert_includes(ref_lines, 1) # def self.bar declaration
-    assert_includes(ref_lines, 13) # Foo.bar call site
     refute_includes(ref_lines, 6) # def bar (Qux) must not appear
+    refute_includes(ref_lines, 11) # def self.bar (Other) must not appear
+    assert_includes(ref_lines, 16) # it.bar included because receiver is unresolved
+    assert_includes(ref_lines, 18) # Foo.bar call site
+    refute_includes(ref_lines, 19) # Other.bar call site is filtered out by its resolved receiver
+  end
 
-    # it.bar must not appear, but we don't currently expose method reference receivers in the API
-    assert_includes(ref_lines, 11)
-
-    # Cursor on `bar` in `it.bar` — the type of the local `it` can't be inferred, so we fall back
-    # to including every `bar` declaration and call site as candidates.
-    refs = find_references(source, { line: 11, character: 3 }, include_declarations: true)
+  def test_falls_back_to_all_candidates_when_call_site_receiver_is_unresolvable
+    refs = find_references(AMBIGUOUS_BAR_SOURCE, { line: 16, character: 3 }, include_declarations: true)
     ref_lines = refs.map { |r| r.range.start.line }
+
     assert_includes(ref_lines, 1) # Foo.bar declaration
     assert_includes(ref_lines, 6) # Qux#bar declaration
-    assert_includes(ref_lines, 11) # it.bar call site
-    assert_includes(ref_lines, 13) # Foo.bar call site
+    assert_includes(ref_lines, 11) # Other.bar declaration
+    assert_includes(ref_lines, 16) # it.bar call site
+    assert_includes(ref_lines, 18) # Foo.bar call site
+    assert_includes(ref_lines, 19) # Other.bar call site
+  end
+
+  def test_method_references_match_through_superclass_chain
+    source = <<~RUBY
+      class Parent
+        def self.bar
+        end
+      end
+
+      class Child < Parent
+      end
+
+      Parent.bar
+      Child.bar
+    RUBY
+
+    refs = find_references(source, { line: 1, character: 11 }, include_declarations: true)
+    ref_lines = refs.map { |r| r.range.start.line }
+    assert_includes(ref_lines, 1) # Parent.bar declaration
+    assert_includes(ref_lines, 8) # Parent.bar call site
+    assert_includes(ref_lines, 9) # Child.bar call site (matched through Child::<Child>'s ancestors)
   end
 
   def test_finds_references_from_def_node