Merge #11235

11235: internal: Record attribute calls on assoc items in TraitData and ImplData r=Veykril a=Veykril

Fixes https://github.com/rust-analyzer/rust-analyzer/issues/11228

Co-authored-by: Lukas Wirth <lukastw97@gmail.com>

5 files changed, 105 insertions, 42 deletions

diff --git a/crates/hir/src/semantics/source_to_def.rs b/crates/hir/src/semantics/source_to_def.rs
index f9c6564d01f..c0e8d0e082a 100644
--- a/crates/hir/src/semantics/source_to_def.rs
+++ b/crates/hir/src/semantics/source_to_def.rs
@@ -316,8 +316,7 @@ impl SourceToDefCtx<'_, '_> {
     }
 
     pub(super) fn macro_to_def(&mut self, src: InFile<ast::Macro>) -> Option<MacroDefId> {
-        let makro =
-            self.dyn_map(src.as_ref()).and_then(|it| it[keys::MACRO_CALL].get(&src).copied());
+        let makro = self.dyn_map(src.as_ref()).and_then(|it| it[keys::MACRO].get(&src).copied());
         if let res @ Some(_) = makro {
             return res;
         }
diff --git a/crates/hir_def/src/child_by_source.rs b/crates/hir_def/src/child_by_source.rs
index 1ef41d90b51..545ae41edf5 100644
--- a/crates/hir_def/src/child_by_source.rs
+++ b/crates/hir_def/src/child_by_source.rs
@@ -30,20 +30,31 @@ pub trait ChildBySource {
 impl ChildBySource for TraitId {
     fn child_by_source_to(&self, db: &dyn DefDatabase, res: &mut DynMap, file_id: HirFileId) {
         let data = db.trait_data(*self);
-        // FIXME attribute macros
-        for &(_, item) in data.items.iter() {
+
+        data.attribute_calls().filter(|(ast_id, _)| ast_id.file_id == file_id).for_each(
+            |(ast_id, call_id)| {
+                let item = ast_id.with_value(ast_id.to_node(db.upcast()));
+                res[keys::ATTR_MACRO_CALL].insert(item, call_id);
+            },
+        );
+        data.items.iter().for_each(|&(_, item)| {
             child_by_source_assoc_items(db, res, file_id, item);
-        }
+        });
     }
 }
 
 impl ChildBySource for ImplId {
     fn child_by_source_to(&self, db: &dyn DefDatabase, res: &mut DynMap, file_id: HirFileId) {
         let data = db.impl_data(*self);
-        // FIXME attribute macros
-        for &item in data.items.iter() {
+        data.attribute_calls().filter(|(ast_id, _)| ast_id.file_id == file_id).for_each(
+            |(ast_id, call_id)| {
+                let item = ast_id.with_value(ast_id.to_node(db.upcast()));
+                res[keys::ATTR_MACRO_CALL].insert(item, call_id);
+            },
+        );
+        data.items.iter().for_each(|&item| {
             child_by_source_assoc_items(db, res, file_id, item);
-        }
+        });
     }
 }
 
@@ -97,7 +108,7 @@ impl ChildBySource for ItemScope {
                     // FIXME: Do we need to add proc-macros into a PROCMACRO dynmap here?
                     Either::Right(_fn) => return,
                 };
-                res[keys::MACRO_CALL].insert(src, makro);
+                res[keys::MACRO].insert(src, makro);
             }
         });
         self.unnamed_consts().for_each(|konst| {
diff --git a/crates/hir_def/src/data.rs b/crates/hir_def/src/data.rs
index 753084fb4bc..471c7330d73 100644
--- a/crates/hir_def/src/data.rs
+++ b/crates/hir_def/src/data.rs
@@ -2,7 +2,7 @@
 
 use std::sync::Arc;
 
-use hir_expand::{name::Name, AstId, ExpandResult, InFile};
+use hir_expand::{name::Name, AstId, ExpandResult, InFile, MacroCallId};
 use syntax::ast;
 
 use crate::{
@@ -184,6 +184,8 @@ pub struct TraitData {
     /// method calls to this trait's methods when the receiver is an array and the crate edition is
     /// 2015 or 2018.
     pub skip_array_during_method_dispatch: bool,
+    // box it as the vec is usually empty anyways
+    pub attribute_calls: Option<Box<Vec<(AstId<ast::Item>, MacroCallId)>>>,
 }
 
 impl TraitData {
@@ -207,15 +209,8 @@ impl TraitData {
             .by_key("rustc_skip_array_during_method_dispatch")
             .exists();
 
-        let items = collect_items(
-            db,
-            module_id,
-            &mut expander,
-            tr_def.items.iter().copied(),
-            tr_loc.id.tree_id(),
-            container,
-            100,
-        );
+        let (items, attribute_calls) =
+            do_collect(db, module_id, &mut expander, &tr_def.items, tr_loc.id.tree_id(), container);
 
         Arc::new(TraitData {
             name,
@@ -224,6 +219,7 @@ impl TraitData {
             is_unsafe,
             visibility,
             skip_array_during_method_dispatch,
+            attribute_calls,
         })
     }
 
@@ -247,6 +243,10 @@ impl TraitData {
             _ => None,
         })
     }
+
+    pub fn attribute_calls(&self) -> impl Iterator<Item = (AstId<ast::Item>, MacroCallId)> + '_ {
+        self.attribute_calls.iter().flat_map(|it| it.iter()).copied()
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -255,6 +255,8 @@ pub struct ImplData {
     pub self_ty: Interned<TypeRef>,
     pub items: Vec<AssocItemId>,
     pub is_negative: bool,
+    // box it as the vec is usually empty anyways
+    pub attribute_calls: Option<Box<Vec<(AstId<ast::Item>, MacroCallId)>>>,
 }
 
 impl ImplData {
@@ -271,18 +273,21 @@ impl ImplData {
         let container = ItemContainerId::ImplId(id);
         let mut expander = Expander::new(db, impl_loc.id.file_id(), module_id);
 
-        let items = collect_items(
+        let (items, attribute_calls) = do_collect(
             db,
             module_id,
             &mut expander,
-            impl_def.items.iter().copied(),
+            &impl_def.items,
             impl_loc.id.tree_id(),
             container,
-            100,
         );
         let items = items.into_iter().map(|(_, item)| item).collect();
 
-        Arc::new(ImplData { target_trait, self_ty, items, is_negative })
+        Arc::new(ImplData { target_trait, self_ty, items, is_negative, attribute_calls })
+    }
+
+    pub fn attribute_calls(&self) -> impl Iterator<Item = (AstId<ast::Item>, MacroCallId)> + '_ {
+        self.attribute_calls.iter().flat_map(|it| it.iter()).copied()
     }
 }
 
@@ -333,17 +338,47 @@ impl StaticData {
     }
 }
 
+fn do_collect(
+    db: &dyn DefDatabase,
+    module_id: ModuleId,
+    expander: &mut Expander,
+    assoc_items: &[AssocItem],
+    tree_id: item_tree::TreeId,
+    container: ItemContainerId,
+) -> (Vec<(Name, AssocItemId)>, Option<Box<Vec<(AstId<ast::Item>, MacroCallId)>>>) {
+    let mut items = Vec::new();
+    let mut attribute_calls = Vec::new();
+
+    collect_items(
+        db,
+        &mut items,
+        &mut attribute_calls,
+        module_id,
+        expander,
+        assoc_items.iter().copied(),
+        tree_id,
+        container,
+        100,
+    );
+
+    let attribute_calls =
+        if attribute_calls.is_empty() { None } else { Some(Box::new(attribute_calls)) };
+    (items, attribute_calls)
+}
+
 fn collect_items(
     db: &dyn DefDatabase,
+    items: &mut Vec<(Name, AssocItemId)>,
+    attr_calls: &mut Vec<(AstId<ast::Item>, MacroCallId)>,
     module: ModuleId,
     expander: &mut Expander,
     assoc_items: impl Iterator<Item = AssocItem>,
     tree_id: item_tree::TreeId,
     container: ItemContainerId,
     limit: usize,
-) -> Vec<(Name, AssocItemId)> {
+) {
     if limit == 0 {
-        return Vec::new();
+        return;
     }
 
     let item_tree = tree_id.item_tree(db);
@@ -351,7 +386,6 @@ fn collect_items(
     let cfg_options = &crate_graph[module.krate].cfg_options;
     let def_map = module.def_map(db);
 
-    let mut items = Vec::new();
     'items: for item in assoc_items {
         let attrs = item_tree.attrs(db, module.krate, ModItem::from(item).into());
         if !attrs.is_cfg_enabled(cfg_options) {
@@ -359,15 +393,15 @@ fn collect_items(
         }
 
         for attr in &*attrs {
-            let ast_id = AstIdWithPath {
-                path: (*attr.path).clone(),
-                ast_id: AstId::new(expander.current_file_id(), item.ast_id(&item_tree).upcast()),
-            };
+            let ast_id = AstId::new(expander.current_file_id(), item.ast_id(&item_tree).upcast());
+            let ast_id_with_path = AstIdWithPath { path: (*attr.path).clone(), ast_id };
+
             if let Ok(ResolvedAttr::Macro(call_id)) =
-                def_map.resolve_attr_macro(db, module.local_id, ast_id, attr)
+                def_map.resolve_attr_macro(db, module.local_id, ast_id_with_path, attr)
             {
+                attr_calls.push((ast_id, call_id));
                 let res = expander.enter_expand_id(db, call_id);
-                items.extend(collect_macro_items(db, module, expander, container, limit, res));
+                collect_macro_items(db, items, attr_calls, module, expander, container, limit, res);
                 continue 'items;
             }
         }
@@ -401,34 +435,32 @@ fn collect_items(
                 let res = expander.enter_expand(db, call);
 
                 if let Ok(res) = res {
-                    items.extend(collect_macro_items(db, module, expander, container, limit, res));
+                    collect_macro_items(
+                        db, items, attr_calls, module, expander, container, limit, res,
+                    );
                 }
             }
         }
     }
-
-    items
 }
 
 fn collect_macro_items(
     db: &dyn DefDatabase,
+    items: &mut Vec<(Name, AssocItemId)>,
+    attr_calls: &mut Vec<(AstId<ast::Item>, MacroCallId)>,
     module: ModuleId,
     expander: &mut Expander,
     container: ItemContainerId,
     limit: usize,
     res: ExpandResult<Option<(Mark, ast::MacroItems)>>,
-) -> Vec<(Name, AssocItemId)> {
+) {
     if let Some((mark, mac)) = res.value {
         let src: InFile<ast::MacroItems> = expander.to_source(mac);
         let tree_id = item_tree::TreeId::new(src.file_id, None);
         let item_tree = tree_id.item_tree(db);
         let iter = item_tree.top_level_items().iter().filter_map(ModItem::as_assoc_item);
-        let items = collect_items(db, module, expander, iter, tree_id, container, limit - 1);
+        collect_items(db, items, attr_calls, module, expander, iter, tree_id, container, limit - 1);
 
         expander.exit(db, mark);
-
-        return items;
     }
-
-    Vec::new()
 }
diff --git a/crates/hir_def/src/keys.rs b/crates/hir_def/src/keys.rs
index 3a9cf6eb812..eaa08a365a1 100644
--- a/crates/hir_def/src/keys.rs
+++ b/crates/hir_def/src/keys.rs
@@ -32,7 +32,7 @@ pub const TYPE_PARAM: Key<ast::TypeParam, TypeParamId> = Key::new();
 pub const LIFETIME_PARAM: Key<ast::LifetimeParam, LifetimeParamId> = Key::new();
 pub const CONST_PARAM: Key<ast::ConstParam, ConstParamId> = Key::new();
 
-pub const MACRO_CALL: Key<ast::Macro, MacroDefId> = Key::new();
+pub const MACRO: Key<ast::Macro, MacroDefId> = Key::new();
 pub const ATTR_MACRO_CALL: Key<ast::Item, MacroCallId> = Key::new();
 pub const DERIVE_MACRO_CALL: Key<ast::Attr, (AttrId, Box<[Option<MacroCallId>]>)> = Key::new();
 
diff --git a/crates/ide/src/references.rs b/crates/ide/src/references.rs
index 949c6dc686b..5e6f0ef6a57 100644
--- a/crates/ide/src/references.rs
+++ b/crates/ide/src/references.rs
@@ -1514,4 +1514,25 @@ fn func$0() {
             "#]],
         )
     }
+
+    #[test]
+    fn attr_assoc_item() {
+        check(
+            r#"
+//- proc_macros: identity
+
+trait Trait {
+    #[proc_macros::identity]
+    fn func() {
+        Self::func$0();
+    }
+}
+"#,
+            expect![[r#"
+                func Function FileId(0) 48..87 51..55
+
+                FileId(0) 74..78
+            "#]],
+        )
+    }
 }