feat: Downmap tokens inside derive helpers

4 files changed, 51 insertions, 22 deletions

diff --git a/crates/hir-def/src/item_scope.rs b/crates/hir-def/src/item_scope.rs
index b98b2855cb0..579f803ea19 100644
--- 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/nameres.rs b/crates/hir-def/src/nameres.rs
index 3949fbb6e7b..8165055e092 100644
--- a/crates/hir-def/src/nameres.rs
+++ b/crates/hir-def/src/nameres.rs
@@ -57,10 +57,10 @@ 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;
@@ -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, 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,12 +298,19 @@ 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, 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
     }
@@ -307,6 +318,7 @@ impl DefMap {
     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 +479,7 @@ impl DefMap {
             registered_attrs,
             registered_tools,
             fn_proc_macro_mapping,
+            derive_helpers_in_scope,
             proc_macro_loading_error: _,
             block: _,
             edition: _,
@@ -483,6 +496,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
index 67651e06413..e14d29952dd 100644
--- a/crates/hir-def/src/nameres/collector.rs
+++ b/crates/hir-def/src/nameres/collector.rs
@@ -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<'_> {
@@ -1132,8 +1128,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.
@@ -1322,10 +1318,11 @@ impl DefCollector<'_> {
             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
+                    self.def_map
+                        .derive_helpers_in_scope
                         .entry(ast_id.map(|it| it.upcast()))
                         .or_default()
-                        .extend(helpers.iter().cloned());
+                        .extend(helpers.iter().cloned().zip(std::iter::repeat(macro_call_id)));
                 }
             }
         }
@@ -2140,7 +2137,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();
diff --git a/crates/hir-def/src/resolver.rs b/crates/hir-def/src/resolver.rs
index c8d3052102f..0c1ae8a064a 100644
--- a/crates/hir-def/src/resolver.rs
+++ b/crates/hir-def/src/resolver.rs
@@ -448,10 +448,14 @@ impl Resolver {
     }
 
     pub fn krate(&self) -> CrateId {
+        self.def_map().krate()
+    }
+
+    pub fn def_map(&self) -> &DefMap {
         self.scopes
             .get(0)
             .and_then(|scope| match scope {
-                Scope::ModuleScope(m) => Some(m.def_map.krate()),
+                Scope::ModuleScope(m) => Some(&m.def_map),
                 _ => None,
             })
             .expect("module scope invariant violated")