feat: Downmap tokens inside derive helpers

9 files changed, 151 insertions, 75 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")
diff --git a/crates/hir/src/semantics.rs b/crates/hir/src/semantics.rs
index 043f2b7c24d..218e15989ad 100644
--- a/crates/hir/src/semantics.rs
+++ b/crates/hir/src/semantics.rs
@@ -733,6 +733,8 @@ impl<'db> SemanticsImpl<'db> {
             Some(it) => it,
             None => return,
         };
+        let def_map = sa.resolver.def_map();
+
         let mut stack: SmallVec<[_; 4]> = smallvec![InFile::new(sa.file_id, token)];
         let mut cache = self.expansion_info_cache.borrow_mut();
         let mut mcache = self.macro_call_cache.borrow_mut();
@@ -764,7 +766,7 @@ impl<'db> SemanticsImpl<'db> {
         while let Some(token) = stack.pop() {
             self.db.unwind_if_cancelled();
             let was_not_remapped = (|| {
-                // are we inside an attribute macro call
+                // First expand into attribute invocations
                 let containing_attribute_macro_call = self.with_ctx(|ctx| {
                     token.value.parent_ancestors().filter_map(ast::Item::cast).find_map(|item| {
                         if item.attrs().next().is_none() {
@@ -784,53 +786,19 @@ impl<'db> SemanticsImpl<'db> {
                     );
                 }
 
-                // or are we inside a function-like macro call
-                if let Some(tt) =
-                    // FIXME replace map.while_some with take_while once stable
-                    token
-                        .value
-                        .parent_ancestors()
-                        .map(ast::TokenTree::cast)
-                        .while_some()
-                        .last()
-                {
-                    let parent = tt.syntax().parent()?;
-                    // check for derive attribute here
-                    let macro_call = match_ast! {
-                        match parent {
-                            ast::MacroCall(mcall) => mcall,
-                            // attribute we failed expansion for earlier, this might be a derive invocation
-                            // so try downmapping the token into the pseudo derive expansion
-                            // see [hir_expand::builtin_attr_macro] for how the pseudo derive expansion works
-                            ast::Meta(meta) => {
-                                let attr = meta.parent_attr()?;
-                                let adt = attr.syntax().parent().and_then(ast::Adt::cast)?;
-                                let call_id = self.with_ctx(|ctx| {
-                                    let (_, call_id, _) = ctx.attr_to_derive_macro_call(
-                                        token.with_value(&adt),
-                                        token.with_value(attr),
-                                    )?;
-                                    Some(call_id)
-                                })?;
-                                let file_id = call_id.as_file();
-                                return process_expansion_for_token(
-                                    &mut stack,
-                                    file_id,
-                                    Some(adt.into()),
-                                    token.as_ref(),
-                                );
-                            },
-                            _ => return None,
-                        }
-                    };
+                // Then check for token trees, that means we are either in a function-like macro or
+                // secondary attribute inputs
+                let tt = token.value.parent_ancestors().map_while(ast::TokenTree::cast).last()?;
+                let parent = tt.syntax().parent()?;
 
-                    if tt.left_delimiter_token().map_or(false, |it| it == token.value) {
-                        return None;
-                    }
-                    if tt.right_delimiter_token().map_or(false, |it| it == token.value) {
-                        return None;
-                    }
+                if tt.left_delimiter_token().map_or(false, |it| it == token.value) {
+                    return None;
+                }
+                if tt.right_delimiter_token().map_or(false, |it| it == token.value) {
+                    return None;
+                }
 
+                if let Some(macro_call) = ast::MacroCall::cast(parent.clone()) {
                     let mcall = token.with_value(macro_call);
                     let file_id = match mcache.get(&mcall) {
                         Some(&it) => it,
@@ -840,11 +808,75 @@ impl<'db> SemanticsImpl<'db> {
                             it
                         }
                     };
-                    return process_expansion_for_token(&mut stack, file_id, None, token.as_ref());
+                    process_expansion_for_token(&mut stack, file_id, None, token.as_ref())
+                } else if let Some(meta) = ast::Meta::cast(parent.clone()) {
+                    // attribute we failed expansion for earlier, this might be a derive invocation
+                    // or derive helper attribute
+                    let attr = meta.parent_attr()?;
+
+                    let adt = if let Some(adt) = attr.syntax().parent().and_then(ast::Adt::cast) {
+                        // this might be a derive, or a derive helper on an ADT
+                        let derive_call = self.with_ctx(|ctx| {
+                            // so try downmapping the token into the pseudo derive expansion
+                            // see [hir_expand::builtin_attr_macro] for how the pseudo derive expansion works
+                            ctx.attr_to_derive_macro_call(
+                                token.with_value(&adt),
+                                token.with_value(attr.clone()),
+                            )
+                            .map(|(_, call_id, _)| call_id)
+                        });
+
+                        match derive_call {
+                            Some(call_id) => {
+                                // resolved to a derive
+                                let file_id = call_id.as_file();
+                                return process_expansion_for_token(
+                                    &mut stack,
+                                    file_id,
+                                    Some(adt.into()),
+                                    token.as_ref(),
+                                );
+                            }
+                            None => Some(adt),
+                        }
+                    } else {
+                        // Otherwise this could be a derive helper on a variant or field
+                        if let Some(field) = attr.syntax().parent().and_then(ast::RecordField::cast)
+                        {
+                            field.syntax().ancestors().take(4).find_map(ast::Adt::cast)
+                        } else if let Some(field) =
+                            attr.syntax().parent().and_then(ast::TupleField::cast)
+                        {
+                            field.syntax().ancestors().take(4).find_map(ast::Adt::cast)
+                        } else if let Some(variant) =
+                            attr.syntax().parent().and_then(ast::Variant::cast)
+                        {
+                            variant.syntax().ancestors().nth(2).and_then(ast::Adt::cast)
+                        } else {
+                            None
+                        }
+                    }?;
+
+                    // Not an attribute, nor a derive, so it's either a builtin or a derive helper
+                    // Try to resolve to a derive helper and downmap
+                    let attr_name = attr.path().and_then(|it| it.as_single_name_ref())?.as_name();
+                    let id = self.db.ast_id_map(token.file_id).ast_id(&adt);
+                    let helpers =
+                        def_map.derive_helpers_in_scope(InFile::new(token.file_id, id))?;
+                    let item = Some(adt.into());
+                    let mut res = None;
+                    for (_, derive) in helpers.iter().filter(|(helper, _)| *helper == attr_name) {
+                        res = res.or(process_expansion_for_token(
+                            &mut stack,
+                            derive.as_file(),
+                            item.clone(),
+                            token.as_ref(),
+                        ));
+                    }
+                    res
+                } else {
+                    None
                 }
-
-                // outside of a macro invocation so this is a "final" token
-                None
             })()
             .is_none();
 
diff --git a/crates/hir/src/semantics/source_to_def.rs b/crates/hir/src/semantics/source_to_def.rs
index 5c4cfa7b45a..ba9a1cfb6b5 100644
--- a/crates/hir/src/semantics/source_to_def.rs
+++ b/crates/hir/src/semantics/source_to_def.rs
@@ -247,6 +247,7 @@ impl SourceToDefCtx<'_, '_> {
         map[keys::ATTR_MACRO_CALL].get(&src.value).copied()
     }
 
+    /// (AttrId, derive attribute call id, derive call ids)
     pub(super) fn attr_to_derive_macro_call(
         &mut self,
         item: InFile<&ast::Adt>,
@@ -257,6 +258,7 @@ impl SourceToDefCtx<'_, '_> {
             .get(&src.value)
             .map(|&(attr_id, call_id, ref ids)| (attr_id, call_id, &**ids))
     }
+
     pub(super) fn has_derives(&mut self, adt: InFile<&ast::Adt>) -> bool {
         self.dyn_map(adt).as_ref().map_or(false, |map| !map[keys::DERIVE_MACRO_CALL].is_empty())
     }
diff --git a/crates/ide-db/src/rename.rs b/crates/ide-db/src/rename.rs
index bb466e43e75..dc1626f4396 100644
--- a/crates/ide-db/src/rename.rs
+++ b/crates/ide-db/src/rename.rs
@@ -316,14 +316,20 @@ pub fn source_edit_from_references(
     // macros can cause multiple refs to occur for the same text range, so keep track of what we have edited so far
     let mut edited_ranges = Vec::new();
     for &FileReference { range, ref name, .. } in references {
+        let name_range = name.syntax().text_range();
+        if name_range.len() != range.len() {
+            // This usage comes from a different token kind that was downmapped to a NameLike in a macro
+            // Renaming this will most likely break things syntax-wise
+            continue;
+        }
         let has_emitted_edit = match name {
             // if the ranges differ then the node is inside a macro call, we can't really attempt
             // to make special rewrites like shorthand syntax and such, so just rename the node in
             // the macro input
-            ast::NameLike::NameRef(name_ref) if name_ref.syntax().text_range() == range => {
+            ast::NameLike::NameRef(name_ref) if name_range == range => {
                 source_edit_from_name_ref(&mut edit, name_ref, new_name, def)
             }
-            ast::NameLike::Name(name) if name.syntax().text_range() == range => {
+            ast::NameLike::Name(name) if name_range == range => {
                 source_edit_from_name(&mut edit, name, new_name)
             }
             _ => false,
diff --git a/crates/ide-db/src/search.rs b/crates/ide-db/src/search.rs
index c75364084e3..eb4fc364381 100644
--- a/crates/ide-db/src/search.rs
+++ b/crates/ide-db/src/search.rs
@@ -54,7 +54,9 @@ impl IntoIterator for UsageSearchResult {
 
 #[derive(Debug, Clone)]
 pub struct FileReference {
+    /// The range of the reference in the original file
     pub range: TextRange,
+    /// The node of the reference in the (macro-)file
     pub name: ast::NameLike,
     pub category: Option<ReferenceCategory>,
 }
diff --git a/crates/ide/src/hover.rs b/crates/ide/src/hover.rs
index d8867cf783a..59c97f2dcf9 100644
--- a/crates/ide/src/hover.rs
+++ b/crates/ide/src/hover.rs
@@ -115,7 +115,12 @@ pub(crate) fn hover(
         });
     }
 
-    let descended = sema.descend_into_macros_with_same_text(original_token.clone());
+    let in_attr = matches!(original_token.parent().and_then(ast::TokenTree::cast), Some(tt) if tt.syntax().ancestors().any(|it| ast::Meta::can_cast(it.kind())));
+    let descended = if in_attr {
+        [sema.descend_into_macros_with_kind_preference(original_token.clone())].into()
+    } else {
+        sema.descend_into_macros_with_same_text(original_token.clone())
+    };
 
     // FIXME: Definition should include known lints and the like instead of having this special case here
     let hovered_lint = descended.iter().find_map(|token| {