Expose Declaration#references to the Ruby API (#635)

This PR exposes the `Declaration#references` Ruby API, which is what the LSP will use to implement find references and rename.

Similar to the `declarations` iterator, I extracted the base infrastructure for iterating over references to `utils.c`, so that we can more easily reuse.

Author: vinistock

ext/rubydex/declaration.c

@@ -2,6 +2,7 @@
 #include "definition.h"
 #include "graph.h"
 #include "handle.h"
+#include "reference.h"
 #include "rustbindings.h"
 #include "utils.h"
 
@@ -291,6 +292,42 @@ static VALUE rdxr_declaration_descendants(VALUE self) {
     return self;
 }
 
+// Size function for the Declaration#references enumerator
+static VALUE declaration_references_size(VALUE self, VALUE _args, VALUE _eobj) {
+    HandleData *data;
+    TypedData_Get_Struct(self, HandleData, &handle_type, data);
+
+    void *graph;
+    TypedData_Get_Struct(data->graph_obj, void *, &graph_type, graph);
+
+    struct ReferencesIter *iter = rdx_declaration_references_iter_new(graph, data->id);
+    size_t len = rdx_references_iter_len(iter);
+    rdx_references_iter_free(iter);
+
+    return SIZET2NUM(len);
+}
+
+// Returns an enumerator for all references to this declaration
+//
+// Declaration#references: () -> Enumerator[Reference]
+static VALUE rdxr_declaration_references(VALUE self) {
+    if (!rb_block_given_p()) {
+        return rb_enumeratorize_with_size(self, rb_str_new2("references"), 0, NULL, declaration_references_size);
+    }
+
+    HandleData *data;
+    TypedData_Get_Struct(self, HandleData, &handle_type, data);
+
+    void *graph;
+    TypedData_Get_Struct(data->graph_obj, void *, &graph_type, graph);
+
+    void *iter = rdx_declaration_references_iter_new(graph, data->id);
+    VALUE args = rb_ary_new_from_args(2, data->graph_obj, ULL2NUM((uintptr_t)iter));
+    rb_ensure(rdxi_references_yield, args, rdxi_references_ensure, args);
+
+    return self;
+}
+
 void rdxi_initialize_declaration(VALUE mRubydex) {
     cDeclaration = rb_define_class_under(mRubydex, "Declaration", rb_cObject);
     cNamespace = rb_define_class_under(mRubydex, "Namespace", cDeclaration);
@@ -309,6 +346,7 @@ void rdxi_initialize_declaration(VALUE mRubydex) {
     rb_define_method(cDeclaration, "name", rdxr_declaration_name, 0);
     rb_define_method(cDeclaration, "unqualified_name", rdxr_declaration_unqualified_name, 0);
     rb_define_method(cDeclaration, "definitions", rdxr_declaration_definitions, 0);
+    rb_define_method(cDeclaration, "references", rdxr_declaration_references, 0);
     rb_define_method(cDeclaration, "owner", rdxr_declaration_owner, 0);
 
     // Namespace only methods

ext/rubydex/graph.c

@@ -228,31 +228,6 @@ static VALUE rdxr_graph_aref(VALUE self, VALUE key) {
     return rb_class_new_instance(2, argv, decl_class);
 }
 
-// Body function for rb_ensure for the reference enumerators
-static VALUE graph_references_yield(VALUE args) {
-    VALUE self = rb_ary_entry(args, 0);
-    void *iter = (void *)(uintptr_t)NUM2ULL(rb_ary_entry(args, 1));
-
-    uint64_t id = 0;
-    ReferenceKind kind;
-    while (rdx_references_iter_next(iter, &id, &kind)) {
-        VALUE ref_class = rdxi_reference_class_for_kind(kind);
-        VALUE argv[] = {self, ULL2NUM(id)};
-        VALUE obj = rb_class_new_instance(2, argv, ref_class);
-        rb_yield(obj);
-    }
-
-    return Qnil;
-}
-
-// Ensure function for rb_ensure for the reference enumerators to always free the iterator
-static VALUE graph_references_ensure(VALUE args) {
-    void *iter = (void *)(uintptr_t)NUM2ULL(rb_ary_entry(args, 1));
-    rdx_references_iter_free(iter);
-
-    return Qnil;
-}
-
 // Size function for the constant_references enumerator
 static VALUE graph_constant_references_size(VALUE self, VALUE _args, VALUE _eobj) {
     void *graph;
@@ -278,7 +253,7 @@ static VALUE rdxr_graph_constant_references(VALUE self) {
 
     void *iter = rdx_graph_constant_references_iter_new(graph);
     VALUE args = rb_ary_new_from_args(2, self, ULL2NUM((uintptr_t)iter));
-    rb_ensure(graph_references_yield, args, graph_references_ensure, args);
+    rb_ensure(rdxi_references_yield, args, rdxi_references_ensure, args);
 
     return self;
 }
@@ -308,7 +283,7 @@ static VALUE rdxr_graph_method_references(VALUE self) {
 
     void *iter = rdx_graph_method_references_iter_new(graph);
     VALUE args = rb_ary_new_from_args(2, self, ULL2NUM((uintptr_t)iter));
-    rb_ensure(graph_references_yield, args, graph_references_ensure, args);
+    rb_ensure(rdxi_references_yield, args, rdxi_references_ensure, args);
 
     return self;
 }

ext/rubydex/utils.c

@@ -1,5 +1,6 @@
 #include "utils.h"
 #include "declaration.h"
+#include "reference.h"
 #include "rustbindings.h"
 
 // Convert a Ruby array of strings into a double char pointer so that we can pass that to Rust.
@@ -50,3 +51,28 @@ VALUE rdxi_declarations_ensure(VALUE args) {
     rdx_graph_declarations_iter_free(iter);
     return Qnil;
 }
+
+// Yield body for iterating over references
+VALUE rdxi_references_yield(VALUE args) {
+    VALUE graph_obj = rb_ary_entry(args, 0);
+    void *iter = (void *)(uintptr_t)NUM2ULL(rb_ary_entry(args, 1));
+
+    uint64_t id = 0;
+    ReferenceKind kind;
+
+    while (rdx_references_iter_next(iter, &id, &kind)) {
+        VALUE ref_class = rdxi_reference_class_for_kind(kind);
+        VALUE argv[] = {graph_obj, ULL2NUM(id)};
+        VALUE obj = rb_class_new_instance(2, argv, ref_class);
+        rb_yield(obj);
+    }
+
+    return Qnil;
+}
+
+// Ensure function for iterating over references to always free the iterator
+VALUE rdxi_references_ensure(VALUE args) {
+    void *iter = (void *)(uintptr_t)NUM2ULL(rb_ary_entry(args, 1));
+    rdx_references_iter_free(iter);
+    return Qnil;
+}

ext/rubydex/utils.h

@@ -16,4 +16,10 @@ VALUE rdxi_declarations_yield(VALUE args);
 // Ensure function for iterating over declarations to always free the iterator
 VALUE rdxi_declarations_ensure(VALUE args);
 
+// Yield body for iterating over references
+VALUE rdxi_references_yield(VALUE args);
+
+// Ensure function for iterating over references to always free the iterator
+VALUE rdxi_references_ensure(VALUE args);
+
 #endif // RUBYDEX_UTILS_H

rust/rubydex-sys/src/declaration_api.rs

@@ -7,6 +7,7 @@ use std::ptr;
 
 use crate::definition_api::{DefinitionsIter, rdx_definitions_iter_new_from_ids};
 use crate::graph_api::{GraphPointer, with_graph};
+use crate::reference_api::{ReferenceKind, ReferencesIter};
 use crate::utils;
 use rubydex::model::ids::{DeclarationId, StringId};
 
@@ -399,3 +400,36 @@ pub unsafe extern "C" fn rdx_declaration_descendants(pointer: GraphPointer, decl
 
     Box::into_raw(Box::new(DeclarationsIter::new(declarations.into_boxed_slice())))
 }
+
+/// Creates a new iterator over references for a given declaration by snapshotting the current set of IDs.
+///
+/// # Safety
+///
+/// - `pointer` must be a valid `GraphPointer` previously returned by this crate.
+/// - The returned pointer must be freed with `rdx_references_iter_free`.
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn rdx_declaration_references_iter_new(
+    pointer: GraphPointer,
+    decl_id: u64,
+) -> *mut ReferencesIter {
+    with_graph(pointer, |graph| {
+        let decl_id = DeclarationId::new(decl_id);
+
+        let Some(decl) = graph.declarations().get(&decl_id) else {
+            return ReferencesIter::new(Vec::new().into_boxed_slice());
+        };
+
+        let kind = match decl {
+            Declaration::Namespace(_) | Declaration::Constant(_) | Declaration::ConstantAlias(_) => {
+                ReferenceKind::Constant
+            }
+            Declaration::Method(_) => ReferenceKind::Method,
+            Declaration::GlobalVariable(_) | Declaration::InstanceVariable(_) | Declaration::ClassVariable(_) => {
+                return ReferencesIter::new(Vec::new().into_boxed_slice());
+            }
+        };
+
+        let entries: Vec<_> = decl.references().iter().map(|ref_id| (**ref_id, kind)).collect();
+        ReferencesIter::new(entries.into_boxed_slice())
+    })
+}

test/declaration_test.rb

@@ -380,6 +380,76 @@ class Child < Parent; end
     end
   end
 
+  def test_references_enumerator
+    with_context do |context|
+      context.write!("file1.rb", <<~RUBY)
+        class A; end
+        class B < A; end
+        A.new
+      RUBY
+
+      graph = Rubydex::Graph.new
+      graph.index_all(context.glob("**/*.rb"))
+      graph.resolve
+
+      declaration = graph["A"]
+      enumerator = declaration.references
+      assert_equal(2, enumerator.size)
+      assert_equal(2, enumerator.count)
+      assert_equal(2, enumerator.to_a.size)
+    end
+  end
+
+  def test_references_with_block
+    with_context do |context|
+      context.write!("file1.rb", <<~RUBY)
+        class A; end
+        class B < A; end
+        A.new
+      RUBY
+
+      graph = Rubydex::Graph.new
+      graph.index_all(context.glob("**/*.rb"))
+      graph.resolve
+
+      declaration = graph["A"]
+      references = []
+      declaration.references do |ref|
+        references << ref
+      end
+
+      assert_equal(2, references.size)
+
+      references.each do |ref|
+        assert_instance_of(Rubydex::ConstantReference, ref)
+        assert_equal("A", ref.name)
+        refute_nil(ref.location)
+      end
+    end
+  end
+
+  def test_method_references_are_not_associated_with_declaration
+    # This test documents current behavior. We can only determine all method references with type inference, so we
+    # currently do not try to make the connection
+
+    with_context do |context|
+      context.write!("file1.rb", <<~RUBY)
+        class A
+          def self.foo; end
+        end
+
+        A.foo
+      RUBY
+
+      graph = Rubydex::Graph.new
+      graph.index_all(context.glob("**/*.rb"))
+      graph.resolve
+
+      declaration = graph["A::<A>#foo()"]
+      assert_empty(declaration.references.to_a)
+    end
+  end
+
   def test_find_member_returns_inherited_members
     with_context do |context|
       context.write!("file1.rb", <<~RUBY)