rust-lang · bors · Jul 28, 2022 · Jul 20, 2022 · Jul 23, 2022 · Jul 23, 2022
diff --git a/crates/hir-def/src/data.rs b/crates/hir-def/src/data.rs
@@ -12,7 +12,7 @@ use crate::{
     db::DefDatabase,
     intern::Interned,
     item_tree::{self, AssocItem, FnFlags, ItemTree, ItemTreeId, ModItem, Param, TreeId},
-    nameres::{attr_resolution::ResolvedAttr, DefMap},
+    nameres::{attr_resolution::ResolvedAttr, proc_macro::ProcMacroKind, DefMap},
     type_ref::{TraitRef, TypeBound, TypeRef},
     visibility::RawVisibility,
     AssocItemId, AstIdWithPath, ConstId, ConstLoc, FunctionId, FunctionLoc, HasModule, ImplId,
@@ -348,7 +348,8 @@ impl MacroRulesData {
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct ProcMacroData {
     pub name: Name,
-    // FIXME: Record deriver helper here?
+    /// Derive helpers, if this is a derive
+    pub helpers: Option<Box<[Name]>>,
 }
 
 impl ProcMacroData {
@@ -360,17 +361,23 @@ impl ProcMacroData {
         let item_tree = loc.id.item_tree(db);
         let makro = &item_tree[loc.id.value];
 
-        let name = if let Some(def) = item_tree
+        let (name, helpers) = if let Some(def) = item_tree
             .attrs(db, loc.container.krate(), ModItem::from(loc.id.value).into())
             .parse_proc_macro_decl(&makro.name)
         {
-            def.name
+            (
+                def.name,
+                match def.kind {
+                    ProcMacroKind::CustomDerive { helpers } => Some(helpers),
+                    ProcMacroKind::FnLike | ProcMacroKind::Attr => None,
+                },
+            )
         } else {
             // eeeh...
             stdx::never!("proc macro declaration is not a proc macro");
-            makro.name.clone()
+            (makro.name.clone(), None)
         };
-        Arc::new(ProcMacroData { name })
+        Arc::new(ProcMacroData { name, helpers })
     }
 }
 

diff --git a/crates/hir-def/src/item_scope.rs b/crates/hir-def/src/item_scope.rs
@@ -66,10 +66,14 @@ pub struct ItemScope {
     attr_macros: FxHashMap<AstId<ast::Item>, MacroCallId>,
     /// The derive macro invocations in this scope, keyed by the owner item over the actual derive attributes
     /// paired with the derive macro invocations for the specific attribute.
-    derive_macros: FxHashMap<
-        AstId<ast::Adt>,
-        SmallVec<[(AttrId, MacroCallId, SmallVec<[Option<MacroCallId>; 1]>); 1]>,
-    >,
+    derive_macros: FxHashMap<AstId<ast::Adt>, SmallVec<[DeriveMacroInvocation; 1]>>,
+}
+
+#[derive(Debug, PartialEq, Eq)]
+struct DeriveMacroInvocation {
+    attr_id: AttrId,
+    attr_call_id: MacroCallId,
+    derive_call_ids: SmallVec<[Option<MacroCallId>; 1]>,
 }
 
 pub(crate) static BUILTIN_SCOPE: Lazy<FxHashMap<Name, PerNs>> = Lazy::new(|| {
@@ -210,12 +214,14 @@ impl ItemScope {
         &mut self,
         adt: AstId<ast::Adt>,
         call: MacroCallId,
-        attr_id: AttrId,
+        id: AttrId,
         idx: usize,
     ) {
         if let Some(derives) = self.derive_macros.get_mut(&adt) {
-            if let Some((.., invocs)) = derives.iter_mut().find(|&&mut (id, ..)| id == attr_id) {
-                invocs[idx] = Some(call);
+            if let Some(DeriveMacroInvocation { derive_call_ids, .. }) =
+                derives.iter_mut().find(|&&mut DeriveMacroInvocation { attr_id, .. }| id == attr_id)
+            {
+                derive_call_ids[idx] = Some(call);
             }
         }
     }
@@ -227,10 +233,14 @@ impl ItemScope {
         &mut self,
         adt: AstId<ast::Adt>,
         attr_id: AttrId,
-        call_id: MacroCallId,
+        attr_call_id: MacroCallId,
         len: usize,
     ) {
-        self.derive_macros.entry(adt).or_default().push((attr_id, call_id, smallvec![None; len]));
+        self.derive_macros.entry(adt).or_default().push(DeriveMacroInvocation {
+            attr_id,
+            attr_call_id,
+            derive_call_ids: smallvec![None; len],
+        });
     }
 
     pub(crate) fn derive_macro_invocs(
@@ -242,7 +252,12 @@ impl ItemScope {
         ),
     > + '_ {
         self.derive_macros.iter().map(|(k, v)| {
-            (*k, v.iter().map(|&(attr_id, call_id, ref invocs)| (attr_id, call_id, &**invocs)))
+            (
+                *k,
+                v.iter().map(|DeriveMacroInvocation { attr_id, attr_call_id, derive_call_ids }| {
+                    (*attr_id, *attr_call_id, &**derive_call_ids)
+                }),
+            )
         })
     }
 

diff --git a/crates/hir-def/src/lib.rs b/crates/hir-def/src/lib.rs
@@ -934,11 +934,11 @@ fn derive_macro_as_call_id(
     derive_attr: AttrId,
     derive_pos: u32,
     krate: CrateId,
-    resolver: impl Fn(path::ModPath) -> Option<MacroDefId>,
-) -> Result<MacroCallId, UnresolvedMacro> {
-    let def: MacroDefId = resolver(item_attr.path.clone())
+    resolver: impl Fn(path::ModPath) -> Option<(MacroId, MacroDefId)>,
+) -> Result<(MacroId, MacroDefId, MacroCallId), UnresolvedMacro> {
+    let (macro_id, def_id) = resolver(item_attr.path.clone())
         .ok_or_else(|| UnresolvedMacro { path: item_attr.path.clone() })?;
-    let res = def.as_lazy_macro(
+    let call_id = def_id.as_lazy_macro(
         db.upcast(),
         krate,
         MacroCallKind::Derive {
@@ -947,7 +947,7 @@ fn derive_macro_as_call_id(
             derive_attr_index: derive_attr.ast_index,
         },
     );
-    Ok(res)
+    Ok((macro_id, def_id, call_id))
 }
 
 fn attr_macro_as_call_id(

diff --git a/crates/hir-def/src/nameres.rs b/crates/hir-def/src/nameres.rs
@@ -48,19 +48,19 @@
 //! the result
 
 pub mod attr_resolution;
-mod collector;
+pub mod proc_macro;
 pub mod diagnostics;
+mod collector;
 mod mod_resolution;
 mod path_resolution;
-mod proc_macro;
 
 #[cfg(test)]
 mod tests;
 
-use std::{cmp::Ord, sync::Arc};
+use std::{ops::Deref, sync::Arc};
 
 use base_db::{CrateId, Edition, FileId};
-use hir_expand::{name::Name, InFile, MacroDefId};
+use hir_expand::{name::Name, InFile, MacroCallId, MacroDefId};
 use itertools::Itertools;
 use la_arena::Arena;
 use profile::Count;
@@ -76,7 +76,7 @@ use crate::{
     path::ModPath,
     per_ns::PerNs,
     visibility::Visibility,
-    AstId, BlockId, BlockLoc, FunctionId, LocalModuleId, ModuleId, ProcMacroId,
+    AstId, BlockId, BlockLoc, FunctionId, LocalModuleId, MacroId, ModuleId, ProcMacroId,
 };
 
 /// Contains the results of (early) name resolution.
@@ -106,6 +106,9 @@ pub struct DefMap {
     fn_proc_macro_mapping: FxHashMap<FunctionId, ProcMacroId>,
     /// The error that occurred when failing to load the proc-macro dll.
     proc_macro_loading_error: Option<Box<str>>,
+    /// 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, MacroId, MacroCallId)>>,
 
     /// Custom attributes registered with `#![register_attr]`.
     registered_attrs: Vec<SmolStr>,
@@ -275,6 +278,7 @@ impl DefMap {
             exported_derives: FxHashMap::default(),
             fn_proc_macro_mapping: FxHashMap::default(),
             proc_macro_loading_error: None,
+            derive_helpers_in_scope: FxHashMap::default(),
             prelude: None,
             root,
             modules,
@@ -294,19 +298,30 @@ impl DefMap {
     pub fn modules(&self) -> impl Iterator<Item = (LocalModuleId, &ModuleData)> + '_ {
         self.modules.iter()
     }
+
+    pub fn derive_helpers_in_scope(
+        &self,
+        id: AstId<ast::Adt>,
+    ) -> Option<&[(Name, MacroId, MacroCallId)]> {
+        self.derive_helpers_in_scope.get(&id.map(|it| it.upcast())).map(Deref::deref)
+    }
+
     pub fn registered_tools(&self) -> &[SmolStr] {
         &self.registered_tools
     }
+
     pub fn registered_attrs(&self) -> &[SmolStr] {
         &self.registered_attrs
     }
+
     pub fn root(&self) -> LocalModuleId {
         self.root
     }
 
     pub fn fn_as_proc_macro(&self, id: FunctionId) -> Option<ProcMacroId> {
         self.fn_proc_macro_mapping.get(&id).copied()
     }
+
     pub fn proc_macro_loading_error(&self) -> Option<&str> {
         self.proc_macro_loading_error.as_deref()
     }
@@ -467,6 +482,7 @@ impl DefMap {
             registered_attrs,
             registered_tools,
             fn_proc_macro_mapping,
+            derive_helpers_in_scope,
             proc_macro_loading_error: _,
             block: _,
             edition: _,
@@ -483,6 +499,7 @@ impl DefMap {
         registered_attrs.shrink_to_fit();
         registered_tools.shrink_to_fit();
         fn_proc_macro_mapping.shrink_to_fit();
+        derive_helpers_in_scope.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/collector.rs b/crates/hir-def/src/nameres/collector.rs
@@ -18,7 +18,7 @@ use hir_expand::{
     ExpandTo, HirFileId, InFile, MacroCallId, MacroCallKind, MacroCallLoc, MacroDefId,
     MacroDefKind,
 };
-use itertools::Itertools;
+use itertools::{izip, Itertools};
 use la_arena::Idx;
 use limit::Limit;
 use rustc_hash::{FxHashMap, FxHashSet};
@@ -110,7 +110,6 @@ pub(super) fn collect_defs(db: &dyn DefDatabase, mut def_map: DefMap, tree_id: T
         proc_macros,
         from_glob_import: Default::default(),
         skip_attrs: Default::default(),
-        derive_helpers_in_scope: Default::default(),
         is_proc_macro,
     };
     if tree_id.is_block() {
@@ -258,9 +257,6 @@ struct DefCollector<'a> {
     /// This also stores the attributes to skip when we resolve derive helpers and non-macro
     /// non-builtin attributes in general.
     skip_attrs: FxHashMap<InFile<ModItem>, AttrId>,
-    /// 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>>,
 }
 
 impl DefCollector<'_> {
@@ -1059,16 +1055,20 @@ impl DefCollector<'_> {
         };
         let mut res = ReachedFixedPoint::Yes;
         macros.retain(|directive| {
-            let resolver = |path| {
+            let resolver2 = |path| {
                 let resolved_res = self.def_map.resolve_path_fp_with_macro(
                     self.db,
                     ResolveMode::Other,
                     directive.module_id,
                     &path,
                     BuiltinShadowMode::Module,
                 );
-                resolved_res.resolved_def.take_macros().map(|it| macro_id_to_def_id(self.db, it))
+                resolved_res
+                    .resolved_def
+                    .take_macros()
+                    .map(|it| (it, macro_id_to_def_id(self.db, it)))
             };
+            let resolver = |path| resolver2(path).map(|(_, it)| it);
 
             match &directive.kind {
                 MacroDirectiveKind::FnLike { ast_id, expand_to } => {
@@ -1087,21 +1087,37 @@ impl DefCollector<'_> {
                     }
                 }
                 MacroDirectiveKind::Derive { ast_id, derive_attr, derive_pos } => {
-                    let call_id = derive_macro_as_call_id(
+                    let id = derive_macro_as_call_id(
                         self.db,
                         ast_id,
                         *derive_attr,
                         *derive_pos as u32,
                         self.def_map.krate,
-                        &resolver,
+                        &resolver2,
                     );
-                    if let Ok(call_id) = call_id {
+
+                    if let Ok((macro_id, def_id, call_id)) = id {
                         self.def_map.modules[directive.module_id].scope.set_derive_macro_invoc(
                             ast_id.ast_id,
                             call_id,
                             *derive_attr,
                             *derive_pos,
                         );
+                        // Record its helper attributes.
+                        if def_id.krate != self.def_map.krate {
+                            let def_map = self.db.crate_def_map(def_id.krate);
+                            if let Some(helpers) = def_map.exported_derives.get(&def_id) {
+                                self.def_map
+                                    .derive_helpers_in_scope
+                                    .entry(ast_id.ast_id.map(|it| it.upcast()))
+                                    .or_default()
+                                    .extend(izip!(
+                                        helpers.iter().cloned(),
+                                        iter::repeat(macro_id),
+                                        iter::repeat(call_id),
+                                    ));
+                            }
+                        }
 
                         push_resolved(directive, call_id);
                         res = ReachedFixedPoint::No;
@@ -1132,8 +1148,8 @@ impl DefCollector<'_> {
                     };
 
                     if let Some(ident) = path.as_ident() {
-                        if let Some(helpers) = self.derive_helpers_in_scope.get(&ast_id) {
-                            if helpers.contains(ident) {
+                        if let Some(helpers) = self.def_map.derive_helpers_in_scope.get(&ast_id) {
+                            if helpers.iter().any(|(it, ..)| it == ident) {
                                 cov_mark::hit!(resolved_derive_helper);
                                 // Resolved to derive helper. Collect the item's attributes again,
                                 // starting after the derive helper.
@@ -1148,7 +1164,7 @@ impl DefCollector<'_> {
                     };
                     if matches!(
                         def,
-                        MacroDefId {  kind:MacroDefKind::BuiltInAttr(expander, _),.. }
+                        MacroDefId { kind:MacroDefKind::BuiltInAttr(expander, _),.. }
                         if expander.is_derive()
                     ) {
                         // Resolved to `#[derive]`
@@ -1317,19 +1333,6 @@ impl DefCollector<'_> {
             self.def_map.diagnostics.push(diag);
         }
 
-        // If we've just resolved a derive, record its helper attributes.
-        if let MacroCallKind::Derive { ast_id, .. } = &loc.kind {
-            if loc.def.krate != self.def_map.krate {
-                let def_map = self.db.crate_def_map(loc.def.krate);
-                if let Some(helpers) = def_map.exported_derives.get(&loc.def) {
-                    self.derive_helpers_in_scope
-                        .entry(ast_id.map(|it| it.upcast()))
-                        .or_default()
-                        .extend(helpers.iter().cloned());
-                }
-            }
-        }
-
         // Then, fetch and process the item tree. This will reuse the expansion result from above.
         let item_tree = self.db.file_item_tree(file_id);
         let mod_dir = self.mod_dirs[&module_id].clone();
@@ -2140,7 +2143,6 @@ mod tests {
             proc_macros: Default::default(),
             from_glob_import: Default::default(),
             skip_attrs: Default::default(),
-            derive_helpers_in_scope: Default::default(),
             is_proc_macro: false,
         };
         collector.seed_with_top_level();