Follow aliases for completion

rust/rubydex/src/model/graph.rs

@@ -457,6 +457,38 @@ impl Graph {
         Some(targets)
     }
 
+    /// Resolves a constant alias chain to the final non-alias declaration.
+    ///
+    /// Returns `None` if the declaration is not a constant alias, the chain is circular, or the chain leads to an
+    /// unresolved name.
+    #[must_use]
+    pub fn resolve_alias(&self, declaration_id: &DeclarationId) -> Option<DeclarationId> {
+        let mut seen = IdentityHashSet::default();
+        let mut current_id = *declaration_id;
+
+        loop {
+            if !seen.insert(current_id) {
+                return None;
+            }
+
+            if let Some(targets) = self.alias_targets(&current_id)
+                && let Some(&first_target) = targets.first()
+            {
+                if matches!(
+                    self.declarations.get(&first_target),
+                    Some(Declaration::ConstantAlias(_))
+                ) {
+                    current_id = first_target;
+                    continue;
+                }
+
+                return Some(first_target);
+            }
+
+            return None;
+        }
+    }
+
     #[must_use]
     pub fn names(&self) -> &IdentityHashMap<NameId, NameRef> {
         &self.names
@@ -1664,4 +1696,51 @@ mod tests {
         assert_eq!(12, context.graph().names.len());
         assert_eq!(41, context.graph().strings.len());
     }
+
+    #[test]
+    fn resolve_alias_follows_chain_to_namespace() {
+        let mut context = GraphTest::new();
+
+        context.index_uri(
+            "file:///foo.rb",
+            "
+            class Original; end
+            Alias1 = Original
+            Alias2 = Alias1
+            ",
+        );
+        context.resolve();
+
+        let target = context.graph().resolve_alias(&DeclarationId::from("Alias2"));
+        assert_eq!(target, Some(DeclarationId::from("Original")));
+    }
+
+    #[test]
+    fn resolve_alias_returns_none_for_circular_aliases() {
+        let mut context = GraphTest::new();
+
+        context.index_uri(
+            "file:///foo.rb",
+            "
+            module Foo
+              A = B
+              B = A
+            end
+            ",
+        );
+        context.resolve();
+
+        assert_eq!(context.graph().resolve_alias(&DeclarationId::from("Foo::A")), None);
+        assert_eq!(context.graph().resolve_alias(&DeclarationId::from("Foo::B")), None);
+    }
+
+    #[test]
+    fn resolve_alias_returns_none_for_non_alias() {
+        let mut context = GraphTest::new();
+
+        context.index_uri("file:///foo.rb", "class Foo; end");
+        context.resolve();
+
+        assert!(context.graph().resolve_alias(&DeclarationId::from("Foo")).is_none());
+    }
 }

rust/rubydex/src/query.rs

@@ -240,15 +240,29 @@ pub fn completion_candidates<'a>(
     }
 }
 
+/// Resolves a declaration ID to a namespace, following constant aliases if necessary.
+fn resolve_to_namespace(graph: &Graph, decl_id: DeclarationId) -> Result<DeclarationId, Box<dyn Error>> {
+    if let Some(Declaration::Namespace(_)) = graph.declarations().get(&decl_id) {
+        return Ok(decl_id);
+    }
+
+    if let Some(target_id) = graph.resolve_alias(&decl_id)
+        && let Some(Declaration::Namespace(_)) = graph.declarations().get(&target_id)
+    {
+        return Ok(target_id);
+    }
+
+    Err(format!("Expected declaration {decl_id:?} to be a namespace or alias to a namespace").into())
+}
+
 /// Collect completion for a namespace access (e.g.: `Foo::`)
 fn namespace_access_completion<'a>(
     graph: &'a Graph,
     namespace_decl_id: DeclarationId,
     mut context: CompletionContext<'a>,
 ) -> Result<Vec<CompletionCandidate>, Box<dyn Error>> {
-    let Some(Declaration::Namespace(namespace)) = graph.declarations().get(&namespace_decl_id) else {
-        return Err(format!("Expected declaration {namespace_decl_id:?} to be a namespace").into());
-    };
+    let resolved_id = resolve_to_namespace(graph, namespace_decl_id)?;
+    let namespace = graph.declarations().get(&resolved_id).unwrap().as_namespace().unwrap();
     let mut candidates = Vec::new();
 
     // Walk ancestors collecting inherited constants, stopping at Object to avoid surfacing top-level constants
@@ -294,9 +308,8 @@ fn method_call_completion<'a>(
     receiver_decl_id: DeclarationId,
     mut context: CompletionContext<'a>,
 ) -> Result<Vec<CompletionCandidate>, Box<dyn Error>> {
-    let Some(Declaration::Namespace(namespace)) = graph.declarations().get(&receiver_decl_id) else {
-        return Err(format!("Expected declaration {receiver_decl_id:?} to be a namespace").into());
-    };
+    let resolved_id = resolve_to_namespace(graph, receiver_decl_id)?;
+    let namespace = graph.declarations().get(&resolved_id).unwrap().as_namespace().unwrap();
     let mut candidates = Vec::new();
 
     for ancestor in namespace.ancestors() {
@@ -1019,6 +1032,68 @@ mod tests {
         );
     }
 
+    #[test]
+    fn namespace_access_completion_follows_constant_alias() {
+        let mut context = GraphTest::new();
+
+        context.index_uri(
+            "file:///foo.rb",
+            "
+            class Original
+              CONST = 1
+              class Nested; end
+
+              class << self
+                def class_method; end
+              end
+            end
+
+            module Foo
+              MyOriginal = Original
+            end
+            ",
+        );
+        context.resolve();
+
+        assert_completion_eq!(
+            context,
+            CompletionReceiver::NamespaceAccess(DeclarationId::from("Foo::MyOriginal")),
+            [
+                "Original::CONST",
+                "Original::Nested",
+                "Original::<Original>#class_method()"
+            ]
+        );
+    }
+
+    #[test]
+    fn namespace_access_completion_follows_chained_constant_alias() {
+        let mut context = GraphTest::new();
+
+        context.index_uri(
+            "file:///foo.rb",
+            "
+            class Original
+              CONST = 1
+
+              class << self
+                def class_method; end
+              end
+            end
+
+            Alias1 = Original
+            Alias2 = Alias1
+            ",
+        );
+        context.resolve();
+
+        assert_completion_eq!(
+            context,
+            CompletionReceiver::NamespaceAccess(DeclarationId::from("Alias2")),
+            ["Original::CONST", "Original::<Original>#class_method()"]
+        );
+    }
+
     #[test]
     fn namespace_access_completion_on_basic_object_subclass() {
         let mut context = GraphTest::new();
@@ -1108,6 +1183,36 @@ mod tests {
         );
     }
 
+    #[test]
+    fn method_call_completion_follows_constant_alias() {
+        let mut context = GraphTest::new();
+
+        context.index_uri(
+            "file:///foo.rb",
+            "
+            class Original
+              def bar; end
+              def baz; end
+
+              class << self
+                def class_method; end
+              end
+            end
+
+            module Foo
+              MyOriginal = Original
+            end
+            ",
+        );
+        context.resolve();
+
+        assert_completion_eq!(
+            context,
+            CompletionReceiver::MethodCall(DeclarationId::from("Foo::MyOriginal")),
+            ["Original#baz()", "Original#bar()"]
+        );
+    }
+
     #[test]
     fn method_call_completion_includes_inherited_methods() {
         let mut context = GraphTest::new();