fix: Record macro calls in signatures in `ChildBySource` impls

4 files changed, 77 insertions, 19 deletions

diff --git a/src/tools/rust-analyzer/crates/hir/src/semantics.rs b/src/tools/rust-analyzer/crates/hir/src/semantics.rs
index ee83d10c0e9..4a2e8e379fb 100644
--- a/src/tools/rust-analyzer/crates/hir/src/semantics.rs
+++ b/src/tools/rust-analyzer/crates/hir/src/semantics.rs
@@ -1948,18 +1948,12 @@ impl<'db> SemanticsImpl<'db> {
             ChildContainer::TraitId(it) => {
                 return Some(SourceAnalyzer::new_generic_def(self.db, it.into(), node, offset));
             }
-            ChildContainer::TraitAliasId(it) => {
-                return Some(SourceAnalyzer::new_generic_def(self.db, it.into(), node, offset));
-            }
             ChildContainer::ImplId(it) => {
                 return Some(SourceAnalyzer::new_generic_def(self.db, it.into(), node, offset));
             }
             ChildContainer::EnumId(it) => {
                 return Some(SourceAnalyzer::new_generic_def(self.db, it.into(), node, offset));
             }
-            ChildContainer::TypeAliasId(it) => {
-                return Some(SourceAnalyzer::new_generic_def(self.db, it.into(), node, offset));
-            }
             ChildContainer::GenericDefId(it) => {
                 return Some(SourceAnalyzer::new_generic_def(self.db, it, node, offset));
             }
diff --git a/src/tools/rust-analyzer/crates/hir/src/semantics/child_by_source.rs b/src/tools/rust-analyzer/crates/hir/src/semantics/child_by_source.rs
index 6accf9b2e9c..a150df70502 100644
--- a/src/tools/rust-analyzer/crates/hir/src/semantics/child_by_source.rs
+++ b/src/tools/rust-analyzer/crates/hir/src/semantics/child_by_source.rs
@@ -17,6 +17,7 @@ use hir_def::{
         DynMap,
         keys::{self, Key},
     },
+    hir::generics::GenericParams,
     item_scope::ItemScope,
     item_tree::ItemTreeNode,
     nameres::DefMap,
@@ -49,6 +50,12 @@ impl ChildBySource for TraitId {
         data.items.iter().for_each(|&(_, item)| {
             add_assoc_item(db, res, file_id, item);
         });
+        let (_, source_map) = db.trait_signature_with_source_map(*self);
+        source_map.expansions().filter(|(ast, _)| ast.file_id == file_id).for_each(
+            |(ast, &exp_id)| {
+                res[keys::MACRO_CALL].insert(ast.value, exp_id);
+            },
+        );
     }
 }
 
@@ -68,6 +75,12 @@ impl ChildBySource for ImplId {
         data.items.iter().for_each(|&(_, item)| {
             add_assoc_item(db, res, file_id, item);
         });
+        let (_, source_map) = db.impl_signature_with_source_map(*self);
+        source_map.expansions().filter(|(ast, _)| ast.file_id == file_id).for_each(
+            |(ast, &exp_id)| {
+                res[keys::MACRO_CALL].insert(ast.value, exp_id);
+            },
+        );
     }
 }
 
@@ -195,6 +208,12 @@ impl ChildBySource for EnumId {
             res[keys::ENUM_VARIANT]
                 .insert(ast_id_map.get(tree[variant.lookup(db).id.value].ast_id), variant);
         });
+        let (_, source_map) = db.enum_signature_with_source_map(*self);
+        source_map.expansions().filter(|(ast, _)| ast.file_id == file_id).for_each(
+            |(ast, &exp_id)| {
+                res[keys::MACRO_CALL].insert(ast.value, exp_id);
+            },
+        );
     }
 }
 
@@ -225,7 +244,8 @@ impl ChildBySource for GenericDefId {
             return;
         }
 
-        let generic_params = db.generic_params(*self);
+        let (generic_params, _, source_map) =
+            GenericParams::generic_params_and_store_and_source_map(db, *self);
         let mut toc_idx_iter = generic_params.iter_type_or_consts().map(|(idx, _)| idx);
         let lts_idx_iter = generic_params.iter_lt().map(|(idx, _)| idx);
 
@@ -253,6 +273,12 @@ impl ChildBySource for GenericDefId {
                 res[keys::LIFETIME_PARAM].insert(AstPtr::new(&ast_param), id);
             }
         }
+
+        source_map.expansions().filter(|(ast, _)| ast.file_id == file_id).for_each(
+            |(ast, &exp_id)| {
+                res[keys::MACRO_CALL].insert(ast.value, exp_id);
+            },
+        );
     }
 }
 
diff --git a/src/tools/rust-analyzer/crates/hir/src/semantics/source_to_def.rs b/src/tools/rust-analyzer/crates/hir/src/semantics/source_to_def.rs
index b5537d0bff0..71ee0f69389 100644
--- a/src/tools/rust-analyzer/crates/hir/src/semantics/source_to_def.rs
+++ b/src/tools/rust-analyzer/crates/hir/src/semantics/source_to_def.rs
@@ -631,14 +631,14 @@ impl SourceToDefCtx<'_, '_> {
             match &item {
                 ast::Item::Module(it) => self.module_to_def(container.with_value(it))?.into(),
                 ast::Item::Trait(it) => self.trait_to_def(container.with_value(it))?.into(),
-                ast::Item::TraitAlias(it) => {
-                    self.trait_alias_to_def(container.with_value(it))?.into()
-                }
                 ast::Item::Impl(it) => self.impl_to_def(container.with_value(it))?.into(),
                 ast::Item::Enum(it) => self.enum_to_def(container.with_value(it))?.into(),
-                ast::Item::TypeAlias(it) => {
-                    self.type_alias_to_def(container.with_value(it))?.into()
-                }
+                ast::Item::TypeAlias(it) => ChildContainer::GenericDefId(
+                    self.type_alias_to_def(container.with_value(it))?.into(),
+                ),
+                ast::Item::TraitAlias(it) => ChildContainer::GenericDefId(
+                    self.trait_alias_to_def(container.with_value(it))?.into(),
+                ),
                 ast::Item::Struct(it) => {
                     let def = self.struct_to_def(container.with_value(it))?;
                     let is_in_body = it.field_list().is_some_and(|it| {
@@ -738,11 +738,9 @@ pub(crate) enum ChildContainer {
     DefWithBodyId(DefWithBodyId),
     ModuleId(ModuleId),
     TraitId(TraitId),
-    TraitAliasId(TraitAliasId),
     ImplId(ImplId),
     EnumId(EnumId),
     VariantId(VariantId),
-    TypeAliasId(TypeAliasId),
     /// XXX: this might be the same def as, for example an `EnumId`. However,
     /// here the children are generic parameters, and not, eg enum variants.
     GenericDefId(GenericDefId),
@@ -751,11 +749,9 @@ impl_from! {
     DefWithBodyId,
     ModuleId,
     TraitId,
-    TraitAliasId,
     ImplId,
     EnumId,
     VariantId,
-    TypeAliasId,
     GenericDefId
     for ChildContainer
 }
@@ -767,11 +763,9 @@ impl ChildContainer {
             ChildContainer::DefWithBodyId(it) => it.child_by_source(db, file_id),
             ChildContainer::ModuleId(it) => it.child_by_source(db, file_id),
             ChildContainer::TraitId(it) => it.child_by_source(db, file_id),
-            ChildContainer::TraitAliasId(_) => DynMap::default(),
             ChildContainer::ImplId(it) => it.child_by_source(db, file_id),
             ChildContainer::EnumId(it) => it.child_by_source(db, file_id),
             ChildContainer::VariantId(it) => it.child_by_source(db, file_id),
-            ChildContainer::TypeAliasId(_) => DynMap::default(),
             ChildContainer::GenericDefId(it) => it.child_by_source(db, file_id),
         }
     }
diff --git a/src/tools/rust-analyzer/crates/ide/src/expand_macro.rs b/src/tools/rust-analyzer/crates/ide/src/expand_macro.rs
index 7c396339c14..1c09bd5d0c0 100644
--- a/src/tools/rust-analyzer/crates/ide/src/expand_macro.rs
+++ b/src/tools/rust-analyzer/crates/ide/src/expand_macro.rs
@@ -800,4 +800,48 @@ foo();
                 foo();"#]],
         );
     }
+
+    #[test]
+    fn works_in_sig() {
+        check(
+            r#"
+macro_rules! foo {
+    () => { u32 };
+}
+fn foo() -> foo$0!() {
+    42
+}
+"#,
+            expect![[r#"
+                foo!
+                u32"#]],
+        );
+        check(
+            r#"
+macro_rules! foo {
+    () => { u32 };
+}
+fn foo(_: foo$0!() ) {}
+"#,
+            expect![[r#"
+                foo!
+                u32"#]],
+        );
+    }
+
+    #[test]
+    fn works_in_generics() {
+        check(
+            r#"
+trait Trait {}
+macro_rules! foo {
+    () => { Trait };
+}
+impl<const C: foo$0!()> Trait for () {}
+"#,
+            expect![[r#"
+                foo!
+                Trait"#]],
+        );
+    }
 }