Merge #11208

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

fixes #11104

bors r+

Co-authored-by: Jonas Schievink <[email protected]>

Author: bors[bot]

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,17 +156,17 @@ 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())
-                })));
+                }));
             }
         };
 
         let node = match T::cast(raw_node) {
             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,
-    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,28 +401,34 @@ 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));
                 }
             }
         }
     }
 
     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>,
-) -> Result<MacroCallId, UnresolvedMacro> {
+    def: MacroDefId,
+) -> 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().ok_or_else(|| UnresolvedMacro { path: attr_path.clone() })?;
+    let last_segment = attr_path.segments().last().expect("empty attribute path");
     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
+    }
+}