Expand attribute macros on impl and trait items

8 files changed, 200 insertions, 73 deletions

diff --git a/crates/hir_def/src/body.rs b/crates/hir_def/src/body.rs
index 2634c52a0b3..505b07cc8ac 100644
--- a/crates/hir_def/src/body.rs
+++ b/crates/hir_def/src/body.rs
@@ -12,7 +12,8 @@ use cfg::{CfgExpr, CfgOptions};
 use drop_bomb::DropBomb;
 use either::Either;
 use hir_expand::{
-    ast_id_map::AstIdMap, hygiene::Hygiene, AstId, ExpandResult, HirFileId, InFile, MacroDefId,
+    ast_id_map::AstIdMap, hygiene::Hygiene, AstId, ExpandError, ExpandResult, HirFileId, InFile,
+    MacroCallId, MacroDefId,
 };
 use la_arena::{Arena, ArenaMap};
 use limit::Limit;
@@ -124,6 +125,23 @@ impl Expander {
             }
         };
 
+        Ok(self.enter_expand_inner(db, call_id, err))
+    }
+
+    pub fn enter_expand_id<T: ast::AstNode>(
+        &mut self,
+        db: &dyn DefDatabase,
+        call_id: MacroCallId,
+    ) -> ExpandResult<Option<(Mark, T)>> {
+        self.enter_expand_inner(db, call_id, None)
+    }
+
+    fn enter_expand_inner<T: ast::AstNode>(
+        &mut self,
+        db: &dyn DefDatabase,
+        call_id: MacroCallId,
+        mut err: Option<ExpandError>,
+    ) -> ExpandResult<Option<(Mark, T)>> {
         if err.is_none() {
             err = db.macro_expand_error(call_id);
         }
@@ -138,9 +156,9 @@ impl Expander {
                     tracing::warn!("no error despite `parse_or_expand` failing");
                 }
 
-                return Ok(ExpandResult::only_err(err.unwrap_or_else(|| {
+                return ExpandResult::only_err(err.unwrap_or_else(|| {
                     mbe::ExpandError::Other("failed to parse macro invocation".into())
-                })));
+                }));
             }
         };
 
@@ -148,7 +166,7 @@ impl Expander {
             Some(it) => it,
             None => {
                 // This can happen without being an error, so only forward previous errors.
-                return Ok(ExpandResult { value: None, err });
+                return ExpandResult { value: None, err };
             }
         };
 
@@ -164,7 +182,7 @@ impl Expander {
         self.current_file_id = file_id;
         self.ast_id_map = db.ast_id_map(file_id);
 
-        Ok(ExpandResult { value: Some((mark, node)), err })
+        ExpandResult { value: Some((mark, node)), err }
     }
 
     pub fn exit(&mut self, db: &dyn DefDatabase, mut mark: Mark) {
diff --git a/crates/hir_def/src/child_by_source.rs b/crates/hir_def/src/child_by_source.rs
index 5ab236aa3d3..7a9e414ece5 100644
--- a/crates/hir_def/src/child_by_source.rs
+++ b/crates/hir_def/src/child_by_source.rs
@@ -30,6 +30,7 @@ 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 (_name, item) in data.items.iter() {
             match *item {
                 AssocItemId::FunctionId(func) => {
@@ -61,6 +62,7 @@ impl ChildBySource for TraitId {
 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() {
             match item {
                 AssocItemId::FunctionId(func) => {
diff --git a/crates/hir_def/src/data.rs b/crates/hir_def/src/data.rs
index aa2844461b6..753084fb4bc 100644
--- a/crates/hir_def/src/data.rs
+++ b/crates/hir_def/src/data.rs
@@ -2,19 +2,20 @@
 
 use std::sync::Arc;
 
-use hir_expand::{name::Name, InFile};
+use hir_expand::{name::Name, AstId, ExpandResult, InFile};
 use syntax::ast;
 
 use crate::{
     attr::Attrs,
-    body::Expander,
+    body::{Expander, Mark},
     db::DefDatabase,
     intern::Interned,
     item_tree::{self, AssocItem, FnFlags, ItemTreeId, ModItem, Param},
+    nameres::attr_resolution::ResolvedAttr,
     type_ref::{TraitRef, TypeBound, TypeRef},
     visibility::RawVisibility,
-    AssocItemId, ConstId, ConstLoc, FunctionId, FunctionLoc, HasModule, ImplId, Intern,
-    ItemContainerId, Lookup, ModuleId, StaticId, TraitId, TypeAliasId, TypeAliasLoc,
+    AssocItemId, AstIdWithPath, ConstId, ConstLoc, FunctionId, FunctionLoc, HasModule, ImplId,
+    Intern, ItemContainerId, Lookup, ModuleId, StaticId, TraitId, TypeAliasId, TypeAliasLoc,
 };
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -348,14 +349,29 @@ fn collect_items(
     let item_tree = tree_id.item_tree(db);
     let crate_graph = db.crate_graph();
     let cfg_options = &crate_graph[module.krate].cfg_options;
+    let def_map = module.def_map(db);
 
     let mut items = Vec::new();
-    for item in assoc_items {
+    '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()),
+            };
+            if let Ok(ResolvedAttr::Macro(call_id)) =
+                def_map.resolve_attr_macro(db, module.local_id, ast_id, attr)
+            {
+                let res = expander.enter_expand_id(db, call_id);
+                items.extend(collect_macro_items(db, module, expander, container, limit, res));
+                continue 'items;
+            }
+        }
+
         match item {
             AssocItem::Function(id) => {
                 let item = &item_tree[id];
@@ -385,24 +401,7 @@ fn collect_items(
                 let res = expander.enter_expand(db, call);
 
                 if let Ok(res) = res {
-                    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);
-                        items.extend(collect_items(
-                            db,
-                            module,
-                            expander,
-                            iter,
-                            tree_id,
-                            container,
-                            limit - 1,
-                        ));
-
-                        expander.exit(db, mark);
-                    }
+                    items.extend(collect_macro_items(db, module, expander, container, limit, res));
                 }
             }
         }
@@ -410,3 +409,26 @@ fn collect_items(
 
     items
 }
+
+fn collect_macro_items(
+    db: &dyn DefDatabase,
+    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);
+
+        expander.exit(db, mark);
+
+        return items;
+    }
+
+    Vec::new()
+}
diff --git a/crates/hir_def/src/item_tree.rs b/crates/hir_def/src/item_tree.rs
index 45b374f338b..c51201f6752 100644
--- a/crates/hir_def/src/item_tree.rs
+++ b/crates/hir_def/src/item_tree.rs
@@ -920,6 +920,17 @@ impl From<AssocItem> for ModItem {
     }
 }
 
+impl AssocItem {
+    pub fn ast_id(self, tree: &ItemTree) -> FileAstId<ast::AssocItem> {
+        match self {
+            AssocItem::Function(id) => tree[id].ast_id.upcast(),
+            AssocItem::TypeAlias(id) => tree[id].ast_id.upcast(),
+            AssocItem::Const(id) => tree[id].ast_id.upcast(),
+            AssocItem::MacroCall(id) => tree[id].ast_id.upcast(),
+        }
+    }
+}
+
 #[derive(Debug, Eq, PartialEq)]
 pub struct Variant {
     pub name: Name,
diff --git a/crates/hir_def/src/nameres.rs b/crates/hir_def/src/nameres.rs
index cf670e5cc17..f793f0e0023 100644
--- a/crates/hir_def/src/nameres.rs
+++ b/crates/hir_def/src/nameres.rs
@@ -47,6 +47,7 @@
 //! path and, upon success, we run macro expansion and "collect module" phase on
 //! the result
 
+pub mod attr_resolution;
 pub mod diagnostics;
 mod collector;
 mod mod_resolution;
@@ -64,7 +65,7 @@ use la_arena::Arena;
 use profile::Count;
 use rustc_hash::FxHashMap;
 use stdx::format_to;
-use syntax::ast;
+use syntax::{ast, SmolStr};
 
 use crate::{
     db::DefDatabase,
@@ -107,6 +108,11 @@ pub struct DefMap {
     /// (the primary purpose is to resolve derive helpers and fetch a proc-macros name)
     exported_proc_macros: FxHashMap<MacroDefId, ProcMacroDef>,
 
+    /// Custom attributes registered with `#![register_attr]`.
+    registered_attrs: Vec<SmolStr>,
+    /// Custom tool modules registered with `#![register_tool]`.
+    registered_tools: Vec<SmolStr>,
+
     edition: Edition,
     diagnostics: Vec<DefDiagnostic>,
 }
@@ -271,6 +277,8 @@ impl DefMap {
             prelude: None,
             root,
             modules,
+            registered_attrs: Vec::new(),
+            registered_tools: Vec::new(),
             diagnostics: Vec::new(),
         }
     }
@@ -443,6 +451,8 @@ impl DefMap {
             extern_prelude,
             diagnostics,
             modules,
+            registered_attrs,
+            registered_tools,
             block: _,
             edition: _,
             krate: _,
@@ -454,6 +464,8 @@ impl DefMap {
         exported_proc_macros.shrink_to_fit();
         diagnostics.shrink_to_fit();
         modules.shrink_to_fit();
+        registered_attrs.shrink_to_fit();
+        registered_tools.shrink_to_fit();
         for (_, module) in modules.iter_mut() {
             module.children.shrink_to_fit();
             module.scope.shrink_to_fit();
diff --git a/crates/hir_def/src/nameres/attr_resolution.rs b/crates/hir_def/src/nameres/attr_resolution.rs
new file mode 100644
index 00000000000..4a7211b5c98
--- /dev/null
+++ b/crates/hir_def/src/nameres/attr_resolution.rs
@@ -0,0 +1,90 @@
+//! Post-nameres attribute resolution.
+
+use hir_expand::MacroCallId;
+use syntax::{ast, SmolStr};
+
+use crate::{
+    attr::Attr,
+    attr_macro_as_call_id, builtin_attr,
+    db::DefDatabase,
+    item_scope::BuiltinShadowMode,
+    nameres::path_resolution::ResolveMode,
+    path::{ModPath, PathKind},
+    AstIdWithPath, LocalModuleId, UnresolvedMacro,
+};
+
+use super::DefMap;
+
+pub enum ResolvedAttr {
+    /// Attribute resolved to an attribute macro.
+    Macro(MacroCallId),
+    /// Attribute resolved to something else that does not require expansion.
+    Other,
+}
+
+impl DefMap {
+    pub(crate) fn resolve_attr_macro(
+        &self,
+        db: &dyn DefDatabase,
+        original_module: LocalModuleId,
+        ast_id: AstIdWithPath<ast::Item>,
+        attr: &Attr,
+    ) -> Result<ResolvedAttr, UnresolvedMacro> {
+        // NB: does not currently work for derive helpers as they aren't recorded in the `DefMap`
+
+        if self.is_builtin_or_registered_attr(&ast_id.path) {
+            return Ok(ResolvedAttr::Other);
+        }
+
+        let resolved_res = self.resolve_path_fp_with_macro(
+            db,
+            ResolveMode::Other,
+            original_module,
+            &ast_id.path,
+            BuiltinShadowMode::Module,
+        );
+        let def = match resolved_res.resolved_def.take_macros() {
+            Some(def) => {
+                if def.is_attribute() {
+                    def
+                } else {
+                    return Ok(ResolvedAttr::Other);
+                }
+            }
+            None => return Err(UnresolvedMacro { path: ast_id.path.clone() }),
+        };
+
+        Ok(ResolvedAttr::Macro(attr_macro_as_call_id(&ast_id, attr, db, self.krate, def)))
+    }
+
+    pub(crate) fn is_builtin_or_registered_attr(&self, path: &ModPath) -> bool {
+        if path.kind != PathKind::Plain {
+            return false;
+        }
+
+        let segments = path.segments();
+
+        if let Some(name) = segments.first() {
+            let name = name.to_smol_str();
+            let pred = |n: &_| *n == name;
+
+            let registered = self.registered_tools.iter().map(SmolStr::as_str);
+            let is_tool = builtin_attr::TOOL_MODULES.iter().copied().chain(registered).any(pred);
+            // FIXME: tool modules can be shadowed by actual modules
+            if is_tool {
+                return true;
+            }
+
+            if segments.len() == 1 {
+                let registered = self.registered_attrs.iter().map(SmolStr::as_str);
+                let is_inert = builtin_attr::INERT_ATTRIBUTES
+                    .iter()
+                    .map(|it| it.name)
+                    .chain(registered)
+                    .any(pred);
+                return is_inert;
+            }
+        }
+        false
+    }
+}
diff --git a/crates/hir_def/src/nameres/collector.rs b/crates/hir_def/src/nameres/collector.rs
index da774bd0466..e8246d4fb9a 100644
--- a/crates/hir_def/src/nameres/collector.rs
+++ b/crates/hir_def/src/nameres/collector.rs
@@ -20,11 +20,11 @@ use itertools::Itertools;
 use la_arena::Idx;
 use limit::Limit;
 use rustc_hash::{FxHashMap, FxHashSet};
-use syntax::{ast, SmolStr};
+use syntax::ast;
 
 use crate::{
     attr::{Attr, AttrId, AttrInput, Attrs},
-    attr_macro_as_call_id, builtin_attr,
+    attr_macro_as_call_id,
     db::DefDatabase,
     derive_macro_as_call_id,
     intern::Interned,
@@ -97,8 +97,6 @@ pub(super) fn collect_defs(db: &dyn DefDatabase, mut def_map: DefMap, tree_id: T
         from_glob_import: Default::default(),
         skip_attrs: Default::default(),
         derive_helpers_in_scope: Default::default(),
-        registered_attrs: Default::default(),
-        registered_tools: Default::default(),
     };
     if tree_id.is_block() {
         collector.seed_with_inner(tree_id);
@@ -251,10 +249,6 @@ struct DefCollector<'a> {
     /// Tracks which custom derives are in scope for an item, to allow resolution of derive helper
     /// attributes.
     derive_helpers_in_scope: FxHashMap<AstId<ast::Item>, Vec<Name>>,
-    /// Custom attributes registered with `#![register_attr]`.
-    registered_attrs: Vec<SmolStr>,
-    /// Custom tool modules registered with `#![register_tool]`.
-    registered_tools: Vec<SmolStr>,
 }
 
 impl DefCollector<'_> {
@@ -291,10 +285,10 @@ impl DefCollector<'_> {
                 };
 
                 if *attr_name == hir_expand::name![register_attr] {
-                    self.registered_attrs.push(registered_name.to_smol_str());
+                    self.def_map.registered_attrs.push(registered_name.to_smol_str());
                     cov_mark::hit!(register_attr);
                 } else {
-                    self.registered_tools.push(registered_name.to_smol_str());
+                    self.def_map.registered_tools.push(registered_name.to_smol_str());
                     cov_mark::hit!(register_tool);
                 }
             }
@@ -1791,7 +1785,7 @@ impl ModCollector<'_, '_> {
             });
 
         for attr in iter {
-            if self.is_builtin_or_registered_attr(&attr.path) {
+            if self.def_collector.def_map.is_builtin_or_registered_attr(&attr.path) {
                 continue;
             }
             tracing::debug!("non-builtin attribute {}", attr.path);
@@ -1819,37 +1813,6 @@ impl ModCollector<'_, '_> {
         Ok(())
     }
 
-    fn is_builtin_or_registered_attr(&self, path: &ModPath) -> bool {
-        if path.kind != PathKind::Plain {
-            return false;
-        }
-
-        let segments = path.segments();
-
-        if let Some(name) = segments.first() {
-            let name = name.to_smol_str();
-            let pred = |n: &_| *n == name;
-
-            let registered = self.def_collector.registered_tools.iter().map(SmolStr::as_str);
-            let is_tool = builtin_attr::TOOL_MODULES.iter().copied().chain(registered).any(pred);
-            // FIXME: tool modules can be shadowed by actual modules
-            if is_tool {
-                return true;
-            }
-
-            if segments.len() == 1 {
-                let registered = self.def_collector.registered_attrs.iter().map(SmolStr::as_str);
-                let is_inert = builtin_attr::INERT_ATTRIBUTES
-                    .iter()
-                    .map(|it| it.name)
-                    .chain(registered)
-                    .any(pred);
-                return is_inert;
-            }
-        }
-        false
-    }
-
     /// If `attrs` registers a procedural macro, collects its definition.
     fn collect_proc_macro_def(&mut self, func_name: &Name, ast_id: AstId<ast::Fn>, attrs: &Attrs) {
         // FIXME: this should only be done in the root module of `proc-macro` crates, not everywhere
@@ -2101,8 +2064,6 @@ mod tests {
             from_glob_import: Default::default(),
             skip_attrs: Default::default(),
             derive_helpers_in_scope: Default::default(),
-            registered_attrs: Default::default(),
-            registered_tools: Default::default(),
         };
         collector.seed_with_top_level();
         collector.collect();
diff --git a/crates/syntax/src/ast/node_ext.rs b/crates/syntax/src/ast/node_ext.rs
index 98d142da3f7..749581b1ce9 100644
--- a/crates/syntax/src/ast/node_ext.rs
+++ b/crates/syntax/src/ast/node_ext.rs
@@ -108,6 +108,17 @@ impl HasName for Macro {
 
 impl HasAttrs for Macro {}
 
+impl From<ast::AssocItem> for ast::Item {
+    fn from(assoc: ast::AssocItem) -> Self {
+        match assoc {
+            ast::AssocItem::Const(it) => ast::Item::Const(it),
+            ast::AssocItem::Fn(it) => ast::Item::Fn(it),
+            ast::AssocItem::MacroCall(it) => ast::Item::MacroCall(it),
+            ast::AssocItem::TypeAlias(it) => ast::Item::TypeAlias(it),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum AttrKind {
     Inner,