Record attribute calls on assoc items in TraitData and ImplData

5 files changed, 79 insertions, 27 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..f08be39128a 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,7 @@ 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,
+    pub attribute_calls: Option<Box<Vec<(AstId<ast::Item>, MacroCallId)>>>,
 }
 
 impl TraitData {
@@ -207,7 +208,7 @@ impl TraitData {
             .by_key("rustc_skip_array_during_method_dispatch")
             .exists();
 
-        let items = collect_items(
+        let (attribute_calls, items) = collect_items(
             db,
             module_id,
             &mut expander,
@@ -216,6 +217,8 @@ impl TraitData {
             container,
             100,
         );
+        let attribute_calls =
+            if attribute_calls.is_empty() { None } else { Some(Box::new(attribute_calls)) };
 
         Arc::new(TraitData {
             name,
@@ -224,6 +227,7 @@ impl TraitData {
             is_unsafe,
             visibility,
             skip_array_during_method_dispatch,
+            attribute_calls,
         })
     }
 
@@ -247,6 +251,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 +263,7 @@ pub struct ImplData {
     pub self_ty: Interned<TypeRef>,
     pub items: Vec<AssocItemId>,
     pub is_negative: bool,
+    pub attribute_calls: Option<Box<Vec<(AstId<ast::Item>, MacroCallId)>>>,
 }
 
 impl ImplData {
@@ -271,7 +280,7 @@ 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 (attribute_calls, items) = collect_items(
             db,
             module_id,
             &mut expander,
@@ -281,8 +290,14 @@ impl ImplData {
             100,
         );
         let items = items.into_iter().map(|(_, item)| item).collect();
+        let attribute_calls =
+            if attribute_calls.is_empty() { None } else { Some(Box::new(attribute_calls)) };
+
+        Arc::new(ImplData { target_trait, self_ty, items, is_negative, attribute_calls })
+    }
 
-        Arc::new(ImplData { target_trait, self_ty, items, is_negative })
+    pub fn attribute_calls(&self) -> impl Iterator<Item = (AstId<ast::Item>, MacroCallId)> + '_ {
+        self.attribute_calls.iter().flat_map(|it| it.iter()).copied()
     }
 }
 
@@ -341,9 +356,9 @@ fn collect_items(
     tree_id: item_tree::TreeId,
     container: ItemContainerId,
     limit: usize,
-) -> Vec<(Name, AssocItemId)> {
+) -> (Vec<(AstId<ast::Item>, MacroCallId)>, Vec<(Name, AssocItemId)>) {
     if limit == 0 {
-        return Vec::new();
+        return Default::default();
     }
 
     let item_tree = tree_id.item_tree(db);
@@ -352,22 +367,25 @@ fn collect_items(
     let def_map = module.def_map(db);
 
     let mut items = Vec::new();
+    let mut attribute_calls = 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) {
             continue;
         }
-
         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)
             {
+                attribute_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));
+                let (mac_attrs, mac_items) =
+                    collect_macro_items(db, module, expander, container, limit, res);
+                items.extend(mac_items);
+                attribute_calls.extend(mac_attrs);
                 continue 'items;
             }
         }
@@ -401,13 +419,16 @@ 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));
+                    let (mac_attrs, mac_items) =
+                        collect_macro_items(db, module, expander, container, limit, res);
+                    items.extend(mac_items);
+                    attribute_calls.extend(mac_attrs);
                 }
             }
         }
     }
 
-    items
+    (attribute_calls, items)
 }
 
 fn collect_macro_items(
@@ -417,7 +438,7 @@ fn collect_macro_items(
     container: ItemContainerId,
     limit: usize,
     res: ExpandResult<Option<(Mark, ast::MacroItems)>>,
-) -> Vec<(Name, AssocItemId)> {
+) -> (Vec<(AstId<ast::Item>, MacroCallId)>, 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);
@@ -430,5 +451,5 @@ fn collect_macro_items(
         return items;
     }
 
-    Vec::new()
+    Default::default()
 }
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
+            "#]],
+        )
+    }
 }