Specialize reference IDs

rust/rubydex-mcp/src/server.rs

@@ -419,8 +419,16 @@ impl RubydexServer {
         let limit = params.limit.filter(|&l| l > 0).unwrap_or(50).min(200); // default 50, max 200
         let offset = params.offset.unwrap_or(0);
 
+        let empty = rubydex::model::identity_maps::IdentityHashSet::default();
+        let constant_refs = match decl {
+            rubydex::model::declaration::Declaration::Namespace(ns) => ns.references(),
+            rubydex::model::declaration::Declaration::Constant(c) => c.references(),
+            rubydex::model::declaration::Declaration::ConstantAlias(ca) => ca.references(),
+            _ => &empty,
+        };
+
         let (references, total) = paginate!(
-            decl.references().iter(),
+            constant_refs.iter(),
             offset,
             limit,
             |ref_id| {

rust/rubydex-sys/src/declaration_api.rs

@@ -449,21 +449,31 @@ pub unsafe extern "C" fn rdx_declaration_references_iter_new(
             return ReferencesIter::new(Vec::new().into_boxed_slice());
         };
 
-        let kind = match decl {
-            Declaration::Namespace(_) | Declaration::Constant(_) | Declaration::ConstantAlias(_) => {
-                ReferenceKind::Constant
-            }
-            Declaration::Method(_) => ReferenceKind::Method,
+        let entries: Vec<_> = match decl {
+            Declaration::Namespace(ns) => ns
+                .references()
+                .iter()
+                .map(|ref_id| CReference::new(**ref_id, ReferenceKind::Constant))
+                .collect(),
+            Declaration::Constant(c) => c
+                .references()
+                .iter()
+                .map(|ref_id| CReference::new(**ref_id, ReferenceKind::Constant))
+                .collect(),
+            Declaration::ConstantAlias(ca) => ca
+                .references()
+                .iter()
+                .map(|ref_id| CReference::new(**ref_id, ReferenceKind::Constant))
+                .collect(),
+            Declaration::Method(m) => m
+                .references()
+                .iter()
+                .map(|ref_id| CReference::new(**ref_id, ReferenceKind::Method))
+                .collect(),
             Declaration::GlobalVariable(_) | Declaration::InstanceVariable(_) | Declaration::ClassVariable(_) => {
                 return ReferencesIter::new(Vec::new().into_boxed_slice());
             }
         };
-
-        let entries: Vec<_> = decl
-            .references()
-            .iter()
-            .map(|ref_id| CReference::new(**ref_id, kind))
-            .collect();
         ReferencesIter::new(entries.into_boxed_slice())
     })
 }

rust/rubydex-sys/src/reference_api.rs

@@ -5,7 +5,7 @@ use std::ffi::CString;
 use crate::graph_api::{GraphPointer, with_graph};
 use crate::location_api::{Location, create_location_for_uri_and_offset};
 use libc::c_char;
-use rubydex::model::ids::ReferenceId;
+use rubydex::model::ids::{ConstantReferenceId, MethodReferenceId};
 
 /// Kind of reference for FFI dispatch
 #[repr(C)]
@@ -107,7 +107,7 @@ pub unsafe extern "C" fn rdx_references_iter_free(iter: *mut ReferencesIter) {
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn rdx_constant_reference_name(pointer: GraphPointer, reference_id: u64) -> *const c_char {
     with_graph(pointer, |graph| {
-        let ref_id = ReferenceId::new(reference_id);
+        let ref_id = ConstantReferenceId::new(reference_id);
         let reference = graph.constant_references().get(&ref_id).expect("Reference not found");
         let name = graph.names().get(reference.name_id()).expect("Name ID should exist");
 
@@ -133,7 +133,7 @@ pub unsafe extern "C" fn rdx_constant_reference_name(pointer: GraphPointer, refe
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn rdx_method_reference_name(pointer: GraphPointer, reference_id: u64) -> *const c_char {
     with_graph(pointer, |graph| {
-        let ref_id = ReferenceId::new(reference_id);
+        let ref_id = MethodReferenceId::new(reference_id);
         let reference = graph.method_references().get(&ref_id).expect("Reference not found");
         let name = graph
             .strings()
@@ -158,7 +158,7 @@ pub unsafe extern "C" fn rdx_method_reference_name(pointer: GraphPointer, refere
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn rdx_constant_reference_location(pointer: GraphPointer, reference_id: u64) -> *mut Location {
     with_graph(pointer, |graph| {
-        let ref_id = ReferenceId::new(reference_id);
+        let ref_id = ConstantReferenceId::new(reference_id);
         let reference = graph.constant_references().get(&ref_id).expect("Reference not found");
         let document = graph
             .documents()
@@ -183,7 +183,7 @@ pub unsafe extern "C" fn rdx_constant_reference_location(pointer: GraphPointer,
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn rdx_method_reference_location(pointer: GraphPointer, reference_id: u64) -> *mut Location {
     with_graph(pointer, |graph| {
-        let ref_id = ReferenceId::new(reference_id);
+        let ref_id = MethodReferenceId::new(reference_id);
         let reference = graph.method_references().get(&ref_id).expect("Reference not found");
         let document = graph
             .documents()

rust/rubydex/src/indexing/local_graph.rs

@@ -5,7 +5,7 @@ use crate::model::definitions::Definition;
 use crate::model::document::Document;
 use crate::model::graph::NameDependent;
 use crate::model::identity_maps::IdentityHashMap;
-use crate::model::ids::{DefinitionId, NameId, ReferenceId, StringId, UriId};
+use crate::model::ids::{ConstantReferenceId, DefinitionId, MethodReferenceId, NameId, StringId, UriId};
 use crate::model::name::{Name, NameRef};
 use crate::model::references::{ConstantReference, MethodRef};
 use crate::model::string_ref::StringRef;
@@ -17,8 +17,8 @@ type LocalGraphParts = (
     IdentityHashMap<DefinitionId, Definition>,
     IdentityHashMap<StringId, StringRef>,
     IdentityHashMap<NameId, NameRef>,
-    IdentityHashMap<ReferenceId, ConstantReference>,
-    IdentityHashMap<ReferenceId, MethodRef>,
+    IdentityHashMap<ConstantReferenceId, ConstantReference>,
+    IdentityHashMap<MethodReferenceId, MethodRef>,
     IdentityHashMap<NameId, Vec<NameDependent>>,
 );
 
@@ -29,8 +29,8 @@ pub struct LocalGraph {
     definitions: IdentityHashMap<DefinitionId, Definition>,
     strings: IdentityHashMap<StringId, StringRef>,
     names: IdentityHashMap<NameId, NameRef>,
-    constant_references: IdentityHashMap<ReferenceId, ConstantReference>,
-    method_references: IdentityHashMap<ReferenceId, MethodRef>,
+    constant_references: IdentityHashMap<ConstantReferenceId, ConstantReference>,
+    method_references: IdentityHashMap<MethodReferenceId, MethodRef>,
     name_dependents: IdentityHashMap<NameId, Vec<NameDependent>>,
 }
 
@@ -150,11 +150,11 @@ impl LocalGraph {
     // Constant references
 
     #[must_use]
-    pub fn constant_references(&self) -> &IdentityHashMap<ReferenceId, ConstantReference> {
+    pub fn constant_references(&self) -> &IdentityHashMap<ConstantReferenceId, ConstantReference> {
         &self.constant_references
     }
 
-    pub fn add_constant_reference(&mut self, reference: ConstantReference) -> ReferenceId {
+    pub fn add_constant_reference(&mut self, reference: ConstantReference) -> ConstantReferenceId {
         let reference_id = reference.id();
         self.name_dependents
             .entry(*reference.name_id())
@@ -172,11 +172,11 @@ impl LocalGraph {
     // Method references
 
     #[must_use]
-    pub fn method_references(&self) -> &IdentityHashMap<ReferenceId, MethodRef> {
+    pub fn method_references(&self) -> &IdentityHashMap<MethodReferenceId, MethodRef> {
         &self.method_references
     }
 
-    pub fn add_method_reference(&mut self, reference: MethodRef) -> ReferenceId {
+    pub fn add_method_reference(&mut self, reference: MethodRef) -> MethodReferenceId {
         let reference_id = reference.id();
 
         if self.method_references.insert(reference_id, reference).is_some() {

rust/rubydex/src/indexing/rbs_indexer.rs

@@ -16,7 +16,7 @@ use crate::model::definitions::{
     PrependDefinition, Receiver, Signature, Signatures,
 };
 use crate::model::document::Document;
-use crate::model::ids::{DefinitionId, NameId, ReferenceId, StringId, UriId};
+use crate::model::ids::{ConstantReferenceId, DefinitionId, NameId, StringId, UriId};
 use crate::model::name::{Name, ParentScope};
 use crate::model::references::ConstantReference;
 use crate::model::visibility::Visibility;
@@ -124,7 +124,7 @@ impl<'a> RBSIndexer<'a> {
         }
     }
 
-    fn index_mixin(&mut self, type_name: &TypeNameNode, mixin_fn: fn(ReferenceId) -> Mixin) {
+    fn index_mixin(&mut self, type_name: &TypeNameNode, mixin_fn: fn(ConstantReferenceId) -> Mixin) {
         let Some(lexical_nesting_id) = self.parent_lexical_scope_id() else {
             return;
         };

rust/rubydex/src/model/declaration.rs

@@ -1,8 +1,11 @@
 use crate::assert_mem_size;
 use crate::diagnostic::Diagnostic;
+use crate::model::ids::{
+    ClassVariableReferenceId, GlobalVariableReferenceId, InstanceVariableReferenceId, MethodReferenceId,
+};
 use crate::model::{
     identity_maps::{IdentityHashMap, IdentityHashSet},
-    ids::{DeclarationId, DefinitionId, NameId, ReferenceId, StringId},
+    ids::{ConstantReferenceId, DeclarationId, DefinitionId, NameId, StringId},
 };
 
 /// A single ancestor in the linearized ancestor chain
@@ -92,7 +95,7 @@ macro_rules! namespace_declaration {
             /// The list of definition IDs that compose this declaration
             definition_ids: Vec<DefinitionId>,
             /// The set of references that are made to this declaration
-            references: IdentityHashSet<ReferenceId>,
+            references: IdentityHashSet<ConstantReferenceId>,
             /// The ID of the owner of this declaration. For singleton classes, this is the ID of the attached object
             owner_id: DeclarationId,
             /// The entities that are owned by this declaration. For example, constants and methods that are defined inside of
@@ -129,14 +132,32 @@ macro_rules! namespace_declaration {
 
             pub fn extend(&mut self, mut other: Declaration) {
                 self.definition_ids.extend(other.definitions());
-                self.references.extend(other.references());
                 self.diagnostics.extend(other.take_diagnostics());
 
-                if let Declaration::Namespace(namespace) = other {
-                    self.members.extend(namespace.members());
+                match other {
+                    Declaration::Namespace(namespace) => {
+                        self.members.extend(namespace.members());
+                        self.references.extend(namespace.references());
+                    }
+                    Declaration::Constant(constant) => {
+                        self.references.extend(constant.references());
+                    }
+                    Declaration::ConstantAlias(constant_alias) => {
+                        self.references.extend(constant_alias.references());
+                    }
+                    Declaration::Method(_)
+                    | Declaration::GlobalVariable(_)
+                    | Declaration::InstanceVariable(_)
+                    | Declaration::ClassVariable(_) => {
+                        panic!("Cannot extend a namespace declaration with a method declaration");
+                    }
                 }
             }
 
+            pub fn add_reference(&mut self, reference_id: ConstantReferenceId) {
+                self.references.insert(reference_id);
+            }
+
             pub fn set_singleton_class_id(&mut self, declaration_id: DeclarationId) {
                 self.singleton_class_id = Some(declaration_id);
             }
@@ -206,15 +227,15 @@ macro_rules! namespace_declaration {
 
 /// Macro to generate a new struct for simple declarations like variables and methods
 macro_rules! simple_declaration {
-    ($name:ident) => {
+    ($name:ident, $reference_type:ty) => {
         #[derive(Debug)]
         pub struct $name {
             /// The fully qualified name of this declaration
             name: String,
             /// The list of definition IDs that compose this declaration
             definition_ids: Vec<DefinitionId>,
             /// The set of references that are made to this declaration
-            references: IdentityHashSet<ReferenceId>,
+            references: IdentityHashSet<$reference_type>,
             /// The ID of the owner of this declaration
             owner_id: DeclarationId,
             /// Diagnostics associated with this declaration
@@ -233,10 +254,32 @@ macro_rules! simple_declaration {
                 }
             }
 
-            pub fn extend(&mut self, mut other: Declaration) {
+            pub fn extend(&mut self, mut other: $name) {
                 self.definition_ids.extend(other.definitions());
-                self.references.extend(other.references());
                 self.diagnostics.extend(other.take_diagnostics());
+                self.references.extend(other.references());
+            }
+
+            #[must_use]
+            pub fn references(&self) -> &IdentityHashSet<$reference_type> {
+                &self.references
+            }
+
+            pub fn add_reference(&mut self, reference_id: $reference_type) {
+                self.references.insert(reference_id);
+            }
+
+            pub fn remove_reference(&mut self, reference_id: &$reference_type) {
+                self.references.remove(reference_id);
+            }
+
+            pub fn take_diagnostics(&mut self) -> Vec<Diagnostic> {
+                std::mem::take(&mut self.diagnostics)
+            }
+
+            #[must_use]
+            pub fn definitions(&self) -> &[DefinitionId] {
+                &self.definition_ids
             }
         }
     };
@@ -292,11 +335,6 @@ impl Declaration {
         }
     }
 
-    #[must_use]
-    pub fn references(&self) -> &IdentityHashSet<ReferenceId> {
-        all_declarations!(self, it => &it.references)
-    }
-
     #[must_use]
     pub fn definitions(&self) -> &[DefinitionId] {
         all_declarations!(self, it => &it.definition_ids)
@@ -318,18 +356,6 @@ impl Declaration {
         });
     }
 
-    pub fn add_reference(&mut self, id: ReferenceId) {
-        all_declarations!(self, it => {
-            it.references.insert(id);
-        });
-    }
-
-    pub fn remove_reference(&mut self, reference_id: &ReferenceId) {
-        all_declarations!(self, it => {
-            it.references.remove(reference_id);
-        });
-    }
-
     // Deletes a definition from this declaration
     pub fn remove_definition(&mut self, definition_id: &DefinitionId) -> bool {
         all_declarations!(self, it => {
@@ -396,10 +422,22 @@ impl Namespace {
     }
 
     #[must_use]
-    pub fn references(&self) -> &IdentityHashSet<ReferenceId> {
+    pub fn references(&self) -> &IdentityHashSet<ConstantReferenceId> {
         all_namespaces!(self, it => &it.references)
     }
 
+    pub fn remove_reference(&mut self, reference_id: &ConstantReferenceId) {
+        all_namespaces!(self, it => {
+            it.references.remove(reference_id);
+        });
+    }
+
+    pub fn add_reference(&mut self, reference_id: ConstantReferenceId) {
+        all_namespaces!(self, it => {
+            it.references.insert(reference_id);
+        });
+    }
+
     #[must_use]
     pub fn definitions(&self) -> &[DefinitionId] {
         all_namespaces!(self, it => &it.definition_ids)
@@ -514,17 +552,17 @@ namespace_declaration!(SingletonClass, SingletonClassDeclaration);
 assert_mem_size!(SingletonClassDeclaration, 216);
 namespace_declaration!(Todo, TodoDeclaration);
 assert_mem_size!(TodoDeclaration, 216);
-simple_declaration!(ConstantDeclaration);
+simple_declaration!(ConstantDeclaration, ConstantReferenceId);
 assert_mem_size!(ConstantDeclaration, 112);
-simple_declaration!(MethodDeclaration);
+simple_declaration!(MethodDeclaration, MethodReferenceId);
 assert_mem_size!(MethodDeclaration, 112);
-simple_declaration!(GlobalVariableDeclaration);
+simple_declaration!(GlobalVariableDeclaration, GlobalVariableReferenceId);
 assert_mem_size!(GlobalVariableDeclaration, 112);
-simple_declaration!(InstanceVariableDeclaration);
+simple_declaration!(InstanceVariableDeclaration, InstanceVariableReferenceId);
 assert_mem_size!(InstanceVariableDeclaration, 112);
-simple_declaration!(ClassVariableDeclaration);
+simple_declaration!(ClassVariableDeclaration, ClassVariableReferenceId);
 assert_mem_size!(ClassVariableDeclaration, 112);
-simple_declaration!(ConstantAliasDeclaration);
+simple_declaration!(ConstantAliasDeclaration, ConstantReferenceId);
 assert_mem_size!(ConstantAliasDeclaration, 112);
 
 #[cfg(test)]