Merge #11455

11455: Handle proc-macro functions as the proc-macro they resolve to r=Veykril a=Veykril

Fixes https://github.com/rust-analyzer/rust-analyzer/issues/11212

Co-authored-by: Lukas Wirth <lukastw97@gmail.com>

7 files changed, 89 insertions, 9 deletions

diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index a26d8e9ebf9..f047971a116 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -1376,6 +1376,23 @@ impl Function {
         db.function_data(self.id).has_body()
     }
 
+    pub fn as_proc_macro(self, db: &dyn HirDatabase) -> Option<MacroDef> {
+        let function_data = db.function_data(self.id);
+        let attrs = &function_data.attrs;
+        if !(attrs.is_proc_macro()
+            || attrs.is_proc_macro_attribute()
+            || attrs.is_proc_macro_derive())
+        {
+            return None;
+        }
+        let loc = self.id.lookup(db.upcast());
+        let krate = loc.krate(db);
+        let def_map = db.crate_def_map(krate.into());
+        let name = &function_data.name;
+        let mut exported_proc_macros = def_map.exported_proc_macros();
+        exported_proc_macros.find(|(_, mac_name)| mac_name == name).map(|(id, _)| MacroDef { id })
+    }
+
     /// A textual representation of the HIR of this function for debugging purposes.
     pub fn debug_hir(self, db: &dyn HirDatabase) -> String {
         let body = db.body(self.id.into());
diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
index 0a8fb2e47a8..3ff2d7b915e 100644
--- a/crates/hir_def/src/attr.rs
+++ b/crates/hir_def/src/attr.rs
@@ -236,7 +236,9 @@ impl Attrs {
     pub fn by_key(&self, key: &'static str) -> AttrQuery<'_> {
         AttrQuery { attrs: self, key }
     }
+}
 
+impl Attrs {
     pub fn cfg(&self) -> Option<CfgExpr> {
         let mut cfgs = self.by_key("cfg").tt_values().map(CfgExpr::parse).collect::<Vec<_>>();
         match cfgs.len() {
@@ -298,6 +300,18 @@ impl Attrs {
                 matches!(&*tt.token_trees, [tt::TokenTree::Leaf(tt::Leaf::Ident(ident))] if ident.text == "hidden")
         })
     }
+
+    pub fn is_proc_macro(&self) -> bool {
+        self.by_key("proc_macro").exists()
+    }
+
+    pub fn is_proc_macro_attribute(&self) -> bool {
+        self.by_key("proc_macro_attribute").exists()
+    }
+
+    pub fn is_proc_macro_derive(&self) -> bool {
+        self.by_key("proc_macro_derive").exists()
+    }
 }
 
 impl AttrsWithOwner {
diff --git a/crates/hir_def/src/nameres/proc_macro.rs b/crates/hir_def/src/nameres/proc_macro.rs
index 8c3454b31b7..920df7cec20 100644
--- a/crates/hir_def/src/nameres/proc_macro.rs
+++ b/crates/hir_def/src/nameres/proc_macro.rs
@@ -31,12 +31,12 @@ impl ProcMacroKind {
 impl Attrs {
     #[rustfmt::skip]
     pub(super) fn parse_proc_macro_decl(&self, func_name: &Name) -> Option<ProcMacroDef> {
-        if self.by_key("proc_macro").exists() {
+        if self.is_proc_macro() {
             Some(ProcMacroDef { name: func_name.clone(), kind: ProcMacroKind::FnLike })
-        } else if self.by_key("proc_macro_attribute").exists() {
+        } else if self.is_proc_macro_attribute() {
             Some(ProcMacroDef { name: func_name.clone(), kind: ProcMacroKind::Attr })
         } else if self.by_key("proc_macro_derive").exists() {
-            let derive = self.by_key("proc_macro_derive").tt_values().next().unwrap();
+            let derive = self.by_key("proc_macro_derive").tt_values().next()?;
 
             match &*derive.token_trees {
                 // `#[proc_macro_derive(Trait)]`
diff --git a/crates/ide/src/references.rs b/crates/ide/src/references.rs
index 5e6f0ef6a57..91c311fe94e 100644
--- a/crates/ide/src/references.rs
+++ b/crates/ide/src/references.rs
@@ -57,8 +57,6 @@ pub(crate) fn find_all_refs(
     let syntax = sema.parse(position.file_id).syntax().clone();
     let make_searcher = |literal_search: bool| {
         move |def: Definition| {
-            let mut usages =
-                def.usages(sema).set_scope(search_scope.clone()).include_self_refs().all();
             let declaration = match def {
                 Definition::Module(module) => {
                     Some(NavigationTarget::from_module_to_decl(sema.db, module))
@@ -72,6 +70,8 @@ pub(crate) fn find_all_refs(
                     nav,
                 }
             });
+            let mut usages =
+                def.usages(sema).set_scope(search_scope.clone()).include_self_refs().all();
             if literal_search {
                 retain_adt_literal_usages(&mut usages, def, sema);
             }
@@ -1535,4 +1535,47 @@ trait Trait {
             "#]],
         )
     }
+
+    #[test]
+    fn attr() {
+        check(
+            r#"
+//- proc_macros: identity
+
+#[proc_macros::$0identity]
+fn func() {}
+"#,
+            expect![[r#"
+                identity Attribute FileId(1) 1..107 32..40
+
+                FileId(0) 16..24
+            "#]],
+        );
+        check(
+            r#"
+#[proc_macro_attribute]
+fn func$0() {}
+"#,
+            expect![[r#"
+                func Attribute FileId(0) 0..36 27..31
+
+                (no references)
+            "#]],
+        );
+    }
+
+    // FIXME
+    #[test]
+    fn derive() {
+        check(
+            r#"
+//- proc_macros: derive_identity
+//- minicore: derive
+
+#[derive(proc_macros::DeriveIdentity$0)]
+struct Foo;
+"#,
+            expect![[r#""#]],
+        )
+    }
 }
diff --git a/crates/ide_db/src/defs.rs b/crates/ide_db/src/defs.rs
index 172acdbc3c4..5a4cfe6e941 100644
--- a/crates/ide_db/src/defs.rs
+++ b/crates/ide_db/src/defs.rs
@@ -225,7 +225,12 @@ impl NameClass {
                     Definition::Macro(sema.to_def(&ast::Macro::MacroDef(it))?)
                 }
                 ast::Item::Const(it) => Definition::Const(sema.to_def(&it)?),
-                ast::Item::Fn(it) => Definition::Function(sema.to_def(&it)?),
+                ast::Item::Fn(it) => {
+                    let def = sema.to_def(&it)?;
+                    def.as_proc_macro(sema.db)
+                        .map(Definition::Macro)
+                        .unwrap_or(Definition::Function(def))
+                }
                 ast::Item::Module(it) => Definition::Module(sema.to_def(&it)?),
                 ast::Item::Static(it) => Definition::Static(sema.to_def(&it)?),
                 ast::Item::Trait(it) => Definition::Trait(sema.to_def(&it)?),
diff --git a/crates/ide_db/src/helpers.rs b/crates/ide_db/src/helpers.rs
index 9c6d3775c7f..c355016c5df 100644
--- a/crates/ide_db/src/helpers.rs
+++ b/crates/ide_db/src/helpers.rs
@@ -82,6 +82,9 @@ pub fn pick_best_token(
 ) -> Option<SyntaxToken> {
     tokens.max_by_key(move |t| f(t.kind()))
 }
+pub fn pick_token<T: AstToken>(mut tokens: TokenAtOffset<SyntaxToken>) -> Option<T> {
+    tokens.find_map(T::cast)
+}
 
 /// Converts the mod path struct into its ast representation.
 pub fn mod_path_to_ast(path: &hir::ModPath) -> ast::Path {
diff --git a/crates/rust-analyzer/src/handlers.rs b/crates/rust-analyzer/src/handlers.rs
index 4160b3ecd1f..b45fbe698c9 100644
--- a/crates/rust-analyzer/src/handlers.rs
+++ b/crates/rust-analyzer/src/handlers.rs
@@ -275,9 +275,7 @@ pub(crate) fn handle_on_type_formatting(
     let char_typed = params.ch.chars().next().unwrap_or('\0');
 
     let text = snap.analysis.file_text(position.file_id)?;
-    if !text[usize::from(position.offset)..].starts_with(char_typed) {
-        // Add `always!` here once VS Code bug is fixed:
-        //   https://github.com/rust-analyzer/rust-analyzer/issues/10002
+    if stdx::never!(!text[usize::from(position.offset)..].starts_with(char_typed)) {
         return Ok(None);
     }