Merge #11208

11208: feat: expand attribute macros on impl and trait items r=jonas-schievink a=jonas-schievink

fixes https://github.com/rust-analyzer/rust-analyzer/issues/11104

bors r+

Co-authored-by: Jonas Schievink <jonas.schievink@ferrous-systems.com>

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) {

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) => {

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,
+    AssocItemId, AstIdWithPath, ConstId, ConstLoc, FunctionId, FunctionLoc, HasModule, ImplId,
-    ItemContainerId, Lookup, ModuleId, StaticId, TraitId, TypeAliasId, TypeAliasLoc,
+    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 {
+                    items.extend(collect_macro_items(db, module, expander, container, limit, res));
-                        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);
-                    }
                 }
             }
         }
@@ -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()
+}

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,

crates/hir_def/src/lib.rs

@@ -781,12 +781,10 @@ fn attr_macro_as_call_id(
     macro_attr: &Attr,
     db: &dyn db::DefDatabase,
     krate: CrateId,
-    def: Option<MacroDefId>,
+    def: MacroDefId,
-) -> Result<MacroCallId, UnresolvedMacro> {
+) -> MacroCallId {
     let attr_path = &item_attr.path;
-    let def = def.ok_or_else(|| UnresolvedMacro { path: attr_path.clone() })?;
+    let last_segment = attr_path.segments().last().expect("empty attribute path");
-    let last_segment =
-        attr_path.segments().last().ok_or_else(|| UnresolvedMacro { path: attr_path.clone() })?;
     let mut arg = match macro_attr.input.as_deref() {
         Some(attr::AttrInput::TokenTree(tt, map)) => (tt.clone(), map.clone()),
         _ => Default::default(),
@@ -805,5 +803,5 @@ fn attr_macro_as_call_id(
             invoc_attr_index: macro_attr.id.ast_index,
         },
     );
-    Ok(res)
+    res
 }

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();

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
+    }
+}

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);
                 }
             }
@@ -1124,10 +1118,13 @@ impl DefCollector<'_> {
                         }
                     }
 
-                    let def = resolver(path.clone()).filter(MacroDefId::is_attribute);
+                    let def = match resolver(path.clone()) {
+                        Some(def) if def.is_attribute() => def,
+                        _ => return true,
+                    };
                     if matches!(
                         def,
-                        Some(MacroDefId {  kind:MacroDefKind::BuiltInAttr(expander, _),.. })
+                        MacroDefId {  kind:MacroDefKind::BuiltInAttr(expander, _),.. }
                         if expander.is_derive()
                     ) {
                         // Resolved to `#[derive]`
@@ -1184,52 +1181,46 @@ impl DefCollector<'_> {
                         return true;
                     }
 
-                    // Not resolved to a derive helper or the derive attribute, so try to resolve as a normal attribute.
+                    // Not resolved to a derive helper or the derive attribute, so try to treat as a normal attribute.
-                    match attr_macro_as_call_id(file_ast_id, attr, self.db, self.def_map.krate, def)
+                    let call_id =
-                    {
+                        attr_macro_as_call_id(file_ast_id, attr, self.db, self.def_map.krate, def);
-                        Ok(call_id) => {
+                    let loc: MacroCallLoc = self.db.lookup_intern_macro_call(call_id);
-                            let loc: MacroCallLoc = self.db.lookup_intern_macro_call(call_id);
 
-                            // Skip #[test]/#[bench] expansion, which would merely result in more memory usage
+                    // Skip #[test]/#[bench] expansion, which would merely result in more memory usage
-                            // due to duplicating functions into macro expansions
+                    // due to duplicating functions into macro expansions
-                            if matches!(
+                    if matches!(
-                                loc.def.kind,
+                        loc.def.kind,
-                                MacroDefKind::BuiltInAttr(expander, _)
+                        MacroDefKind::BuiltInAttr(expander, _)
-                                if expander.is_test() || expander.is_bench()
+                        if expander.is_test() || expander.is_bench()
-                            ) {
+                    ) {
-                                return recollect_without(self);
+                        return recollect_without(self);
-                            }
-
-                            if let MacroDefKind::ProcMacro(exp, ..) = loc.def.kind {
-                                if exp.is_dummy() {
-                                    // Proc macros that cannot be expanded are treated as not
-                                    // resolved, in order to fall back later.
-                                    self.def_map.diagnostics.push(
-                                        DefDiagnostic::unresolved_proc_macro(
-                                            directive.module_id,
-                                            loc.kind,
-                                        ),
-                                    );
-
-                                    return recollect_without(self);
-                                }
-                            }
-
-                            self.def_map.modules[directive.module_id]
-                                .scope
-                                .add_attr_macro_invoc(ast_id, call_id);
-
-                            resolved.push((
-                                directive.module_id,
-                                call_id,
-                                directive.depth,
-                                directive.container,
-                            ));
-                            res = ReachedFixedPoint::No;
-                            return false;
-                        }
-                        Err(UnresolvedMacro { .. }) => (),
                     }
+
+                    if let MacroDefKind::ProcMacro(exp, ..) = loc.def.kind {
+                        if exp.is_dummy() {
+                            // Proc macros that cannot be expanded are treated as not
+                            // resolved, in order to fall back later.
+                            self.def_map.diagnostics.push(DefDiagnostic::unresolved_proc_macro(
+                                directive.module_id,
+                                loc.kind,
+                            ));
+
+                            return recollect_without(self);
+                        }
+                    }
+
+                    self.def_map.modules[directive.module_id]
+                        .scope
+                        .add_attr_macro_invoc(ast_id, call_id);
+
+                    resolved.push((
+                        directive.module_id,
+                        call_id,
+                        directive.depth,
+                        directive.container,
+                    ));
+                    res = ReachedFixedPoint::No;
+                    return false;
                 }
             }
 
@@ -1794,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);
@@ -1822,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
@@ -2104,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();

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,