Expose method reference receivers

ext/rubydex/reference.c

@@ -75,6 +75,28 @@ static VALUE rdxr_method_reference_location(VALUE self) {
     return location;
 }
 
+// MethodReference#receiver -> Rubydex::Declaration?
+// Returns the resolved declaration for the receiver of the method call. Returns nil when the receiver is not a
+// tracked constant or cannot be resolved.
+static VALUE rdxr_method_reference_receiver(VALUE self) {
+    HandleData *data;
+    TypedData_Get_Struct(self, HandleData, &handle_type, data);
+
+    void *graph;
+    TypedData_Get_Struct(data->graph_obj, void *, &graph_type, graph);
+
+    const struct CDeclaration *decl = rdx_method_reference_receiver_declaration(graph, data->id);
+    if (decl == NULL) {
+        return Qnil;
+    }
+
+    VALUE decl_class = rdxi_declaration_class_for_kind(decl->kind);
+    VALUE argv[] = {data->graph_obj, ULL2NUM(decl->id)};
+    free_c_declaration(decl);
+
+    return rb_class_new_instance(2, argv, decl_class);
+}
+
 // ResolvedConstantReference#declaration -> Declaration
 static VALUE rdxr_resolved_constant_reference_declaration(VALUE self) {
     HandleData *data;
@@ -120,4 +142,5 @@ void rdxi_initialize_reference(VALUE mRubydex) {
     rb_define_method(cMethodReference, "initialize", rdxr_handle_initialize, 2);
     rb_define_method(cMethodReference, "name", rdxr_method_reference_name, 0);
     rb_define_method(cMethodReference, "location", rdxr_method_reference_location, 0);
+    rb_define_method(cMethodReference, "receiver", rdxr_method_reference_receiver, 0);
 }

rbi/rubydex.rbi

@@ -412,6 +412,9 @@ class Rubydex::MethodReference < Rubydex::Reference
 
   sig { returns(String) }
   def name; end
+
+  sig { returns(T.nilable(Rubydex::Declaration)) }
+  def receiver; end
 end
 
 class Rubydex::Reference

rust/rubydex-sys/src/reference_api.rs

@@ -226,6 +226,43 @@ pub unsafe extern "C" fn rdx_resolved_constant_reference_declaration(
     })
 }
 
+/// Returns the declaration of the resolved receiver for the given method reference. Returns NULL when the method
+/// reference has no tracked receiver or when the receiver could not be resolved. Caller must free with
+/// `free_c_declaration`.
+///
+/// # Safety
+///
+/// Assumes pointer is valid.
+///
+/// # Panics
+///
+/// This function will panic if the reference cannot be found.
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn rdx_method_reference_receiver_declaration(
+    pointer: GraphPointer,
+    reference_id: u64,
+) -> *const CDeclaration {
+    with_graph(pointer, |graph| {
+        let ref_id = MethodReferenceId::new(reference_id);
+        let reference = graph.method_references().get(&ref_id).expect("Reference not found");
+
+        let Some(name_id) = reference.receiver() else {
+            return ptr::null();
+        };
+
+        let name_ref = graph.names().get(&name_id).expect("Name ID should exist");
+
+        match name_ref {
+            NameRef::Resolved(resolved) => {
+                let decl_id = *resolved.declaration_id();
+                let decl = graph.declarations().get(&decl_id).expect("Declaration not found");
+                Box::into_raw(Box::new(CDeclaration::from_declaration(decl_id, decl))).cast_const()
+            }
+            NameRef::Unresolved(_) => ptr::null(),
+        }
+    })
+}
+
 /// Returns a newly allocated `Location` for the given method reference id.
 /// Caller must free the returned pointer with `rdx_location_free`.
 ///

test/references_test.rb

@@ -267,4 +267,148 @@ def foo
       assert_equal("#{context.absolute_path_to("file1.rb")}:4:5-4:9", ref1.location.to_display.to_s)
     end
   end
+
+  def test_method_reference_receiver_resolved
+    with_context do |context|
+      context.write!("file1.rb", <<~RUBY)
+        class Foo; end
+        Foo.bar
+      RUBY
+
+      graph = Rubydex::Graph.new
+      graph.index_all(context.glob("**/*.rb"))
+      graph.resolve
+
+      method_ref = graph.method_references.find { |r| r.name == "bar" }
+      refute_nil(method_ref)
+
+      receiver = method_ref.receiver
+      assert_kind_of(Rubydex::SingletonClass, receiver)
+      assert_equal("Foo::<Foo>", receiver.name)
+    end
+  end
+
+  def test_method_reference_receiver_is_nil_when_unresolved
+    with_context do |context|
+      context.write!("file1.rb", <<~RUBY)
+        Bar.baz
+      RUBY
+
+      graph = Rubydex::Graph.new
+      graph.index_all(context.glob("**/*.rb"))
+      graph.resolve
+
+      method_ref = graph.method_references.find { |r| r.name == "baz" }
+      refute_nil(method_ref)
+
+      assert_nil(method_ref.receiver)
+    end
+  end
+
+  def test_method_reference_receiver_implicit_self_in_instance_method
+    with_context do |context|
+      context.write!("file1.rb", <<~RUBY)
+        class Foo
+          def bar
+            baz
+          end
+
+          def baz
+          end
+        end
+      RUBY
+
+      graph = Rubydex::Graph.new
+      graph.index_all(context.glob("**/*.rb"))
+      graph.resolve
+
+      method_ref = graph.method_references.find { |r| r.name == "baz" }
+      refute_nil(method_ref)
+
+      receiver = method_ref.receiver
+      assert_kind_of(Rubydex::Class, receiver)
+      assert_equal("Foo", receiver.name)
+    end
+  end
+
+  def test_method_reference_explicit_self_receiver_in_instance_method
+    with_context do |context|
+      context.write!("file1.rb", <<~RUBY)
+        class Foo
+          def bar
+            self.baz
+          end
+
+          def baz
+          end
+        end
+      RUBY
+
+      graph = Rubydex::Graph.new
+      graph.index_all(context.glob("**/*.rb"))
+      graph.resolve
+
+      method_ref = graph.method_references.find { |r| r.name == "baz" }
+      refute_nil(method_ref)
+
+      receiver = method_ref.receiver
+      assert_kind_of(Rubydex::Class, receiver)
+      assert_equal("Foo", receiver.name)
+    end
+  end
+
+  def test_method_reference_explicit_self_receiver_in_self_method
+    with_context do |context|
+      context.write!("file1.rb", <<~RUBY)
+        class Foo
+          def self.bar
+            self.baz
+          end
+
+          def self.baz
+          end
+        end
+      RUBY
+
+      graph = Rubydex::Graph.new
+      graph.index_all(context.glob("**/*.rb"))
+      graph.resolve
+
+      method_ref = graph.method_references.find { |r| r.name == "baz" }
+      refute_nil(method_ref)
+
+      receiver = method_ref.receiver
+      assert_kind_of(Rubydex::SingletonClass, receiver)
+      assert_equal("<Foo>", receiver.unqualified_name)
+    end
+  end
+
+  def test_method_reference_explicit_self_receiver_in_constant_receiver_method
+    with_context do |context|
+      context.write!("file1.rb", <<~RUBY)
+        class Bar
+          def self.baz
+          end
+        end
+
+        class Foo
+        end
+
+        def Bar.bar
+          self.baz
+        end
+      RUBY
+
+      graph = Rubydex::Graph.new
+      graph.index_all(context.glob("**/*.rb"))
+      graph.resolve
+
+      method_ref = graph.method_references.find { |r| r.name == "baz" }
+      refute_nil(method_ref)
+
+      receiver = method_ref.receiver
+      assert_kind_of(Rubydex::SingletonClass, receiver)
+      assert_equal("Bar::<Bar>", receiver.name)
+    end
+  end
 end