3 files changed, 112 insertions, 32 deletions

diff --git a/crates/hir_def/src/nameres/collector.rs b/crates/hir_def/src/nameres/collector.rs
index d0fefb5af..dcedf7766 100644
--- a/crates/hir_def/src/nameres/collector.rs
+++ b/crates/hir_def/src/nameres/collector.rs
@@ -18,7 +18,6 @@ use hir_expand::{
 use hir_expand::{InFile, MacroCallLoc};
 use rustc_hash::{FxHashMap, FxHashSet};
 use syntax::ast;
-use tt::{Leaf, TokenTree};
 
 use crate::{
     attr::Attrs,
@@ -42,6 +41,8 @@ use crate::{
     UnresolvedMacro,
 };
 
+use super::proc_macro::ProcMacroDef;
+
 const GLOB_RECURSION_LIMIT: usize = 100;
 const EXPANSION_DEPTH_LIMIT: usize = 128;
 const FIXED_POINT_LIMIT: usize = 8192;
@@ -353,24 +354,23 @@ impl DefCollector<'_> {
     /// use a dummy expander that always errors. This comes with the drawback of macros potentially
     /// going out of sync with what the build system sees (since we resolve using VFS state, but
     /// Cargo builds only on-disk files). We could and probably should add diagnostics for that.
-    fn resolve_proc_macro(&mut self, name: &Name) {
+    fn export_proc_macro(&mut self, def: ProcMacroDef, ast_id: AstId<ast::Fn>) {
         self.exports_proc_macros = true;
-        let macro_def = match self.proc_macros.iter().find(|(n, _)| n == name) {
+        let macro_def = match self.proc_macros.iter().find(|(n, _)| n == &def.name) {
             Some((_, expander)) => MacroDefId {
-                ast_id: None,
                 krate: self.def_map.krate,
-                kind: MacroDefKind::ProcMacro(*expander),
+                kind: MacroDefKind::ProcMacro(*expander, ast_id),
                 local_inner: false,
             },
             None => MacroDefId {
-                ast_id: None,
                 krate: self.def_map.krate,
-                kind: MacroDefKind::ProcMacro(ProcMacroExpander::dummy(self.def_map.krate)),
+                kind: MacroDefKind::ProcMacro(ProcMacroExpander::dummy(self.def_map.krate), ast_id),
                 local_inner: false,
             },
         };
 
-        self.define_proc_macro(name.clone(), macro_def);
+        self.define_proc_macro(def.name.clone(), macro_def);
+        self.def_map.exported_proc_macros.insert(macro_def, def);
     }
 
     /// Define a macro with `macro_rules`.
@@ -1118,7 +1118,8 @@ impl ModCollector<'_, '_> {
                 ModItem::Function(id) => {
                     let func = &self.item_tree[id];
 
-                    self.collect_proc_macro_def(&func.name, &attrs);
+                    let ast_id = InFile::new(self.file_id, func.ast_id);
+                    self.collect_proc_macro_def(&func.name, ast_id, &attrs);
 
                     def = Some(DefData {
                         id: FunctionLoc {
@@ -1385,28 +1386,11 @@ impl ModCollector<'_, '_> {
     }
 
     /// If `attrs` registers a procedural macro, collects its definition.
-    fn collect_proc_macro_def(&mut self, func_name: &Name, attrs: &Attrs) {
+    fn collect_proc_macro_def(&mut self, func_name: &Name, ast_id: AstId<ast::Fn>, attrs: &Attrs) {
         // FIXME: this should only be done in the root module of `proc-macro` crates, not everywhere
-        // FIXME: distinguish the type of macro
-        let macro_name = if attrs.by_key("proc_macro").exists()
-            || attrs.by_key("proc_macro_attribute").exists()
-        {
-            func_name.clone()
-        } else {
-            let derive = attrs.by_key("proc_macro_derive");
-            if let Some(arg) = derive.tt_values().next() {
-                if let [TokenTree::Leaf(Leaf::Ident(trait_name)), ..] = &*arg.token_trees {
-                    trait_name.as_name()
-                } else {
-                    log::trace!("malformed `#[proc_macro_derive]`: {}", arg);
-                    return;
-                }
-            } else {
-                return;
-            }
-        };
-
-        self.def_collector.resolve_proc_macro(&macro_name);
+        if let Some(proc_macro) = attrs.parse_proc_macro_decl(func_name) {
+            self.def_collector.export_proc_macro(proc_macro, ast_id);
+        }
     }
 
     fn collect_macro_rules(&mut self, id: FileItemTreeId<MacroRules>) {
@@ -1445,9 +1429,8 @@ impl ModCollector<'_, '_> {
 
         // Case 2: normal `macro_rules!` macro
         let macro_id = MacroDefId {
-            ast_id: Some(ast_id),
             krate: self.def_collector.def_map.krate,
-            kind: MacroDefKind::Declarative,
+            kind: MacroDefKind::Declarative(ast_id),
             local_inner: is_local_inner,
         };
         self.def_collector.define_macro(self.module_id, mac.name.clone(), macro_id, is_export);
diff --git a/crates/hir_def/src/nameres/proc_macro.rs b/crates/hir_def/src/nameres/proc_macro.rs
new file mode 100644
index 000000000..156598f19
--- /dev/null
+++ b/crates/hir_def/src/nameres/proc_macro.rs
@@ -0,0 +1,71 @@
+//! Nameres-specific procedural macro data and helpers.
+
+use hir_expand::name::{AsName, Name};
+use tt::{Leaf, TokenTree};
+
+use crate::attr::Attrs;
+
+#[derive(Debug, PartialEq, Eq)]
+pub(super) struct ProcMacroDef {
+    pub(super) name: Name,
+    pub(super) kind: ProcMacroKind,
+}
+
+#[derive(Debug, PartialEq, Eq)]
+pub(super) enum ProcMacroKind {
+    CustomDerive { helpers: Box<[Name]> },
+    FnLike,
+    Attr,
+}
+
+impl Attrs {
+    #[rustfmt::skip]
+    pub(super) fn parse_proc_macro_decl(&self, func_name: &Name) -> Option<ProcMacroDef> {
+        if self.by_key("proc_macro").exists() {
+            Some(ProcMacroDef { name: func_name.clone(), kind: ProcMacroKind::FnLike })
+        } else if self.by_key("proc_macro_attribute").exists() {
+            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();
+
+            match &*derive.token_trees {
+                // `#[proc_macro_derive(Trait)]`
+                [TokenTree::Leaf(Leaf::Ident(trait_name))] => Some(ProcMacroDef {
+                    name: trait_name.as_name(),
+                    kind: ProcMacroKind::CustomDerive { helpers: Box::new([]) },
+                }),
+
+                // `#[proc_macro_derive(Trait, attibutes(helper1, helper2, ...))]`
+                [
+                    TokenTree::Leaf(Leaf::Ident(trait_name)),
+                    TokenTree::Leaf(Leaf::Punct(comma)),
+                    TokenTree::Leaf(Leaf::Ident(attributes)),
+                    TokenTree::Subtree(helpers)
+                ] if comma.char == ',' && attributes.text == "attributes" =>
+                {
+                    let helpers = helpers.token_trees.iter()
+                        .filter(|tt| !matches!(tt, TokenTree::Leaf(Leaf::Punct(comma)) if comma.char == ','))
+                        .map(|tt| {
+                            match tt {
+                                TokenTree::Leaf(Leaf::Ident(helper)) => Some(helper.as_name()),
+                                _ => None
+                            }
+                        })
+                        .collect::<Option<Box<[_]>>>()?;
+
+                    Some(ProcMacroDef {
+                        name: trait_name.as_name(),
+                        kind: ProcMacroKind::CustomDerive { helpers },
+                    })
+                }
+
+                _ => {
+                    log::trace!("malformed `#[proc_macro_derive]`: {}", derive);
+                    None
+                }
+            }
+        } else {
+            None
+        }
+    }
+}
diff --git a/crates/hir_def/src/nameres/tests/macros.rs b/crates/hir_def/src/nameres/tests/macros.rs
index f65a655bf..d59d3c0db 100644
--- a/crates/hir_def/src/nameres/tests/macros.rs
+++ b/crates/hir_def/src/nameres/tests/macros.rs
@@ -1,4 +1,5 @@
 use super::*;
+use crate::nameres::proc_macro::{ProcMacroDef, ProcMacroKind};
 
 #[test]
 fn macro_rules_are_globally_visible() {
@@ -790,3 +791,28 @@ fn proc_macro_censoring() {
         "#]],
     );
 }
+
+#[test]
+fn collects_derive_helpers() {
+    let def_map = compute_crate_def_map(
+        r"
+        struct TokenStream;
+
+        #[proc_macro_derive(AnotherTrait, attributes(helper_attr))]
+        pub fn derive_macro_2(_item: TokenStream) -> TokenStream {
+            TokenStream
+        }
+        ",
+    );
+
+    assert_eq!(def_map.exported_proc_macros.len(), 1);
+    match def_map.exported_proc_macros.values().next() {
+        Some(ProcMacroDef { kind: ProcMacroKind::CustomDerive { helpers }, .. }) => {
+            match &**helpers {
+                [attr] => assert_eq!(attr.to_string(), "helper_attr"),
+                _ => unreachable!(),
+            }
+        }
+        _ => unreachable!(),
+    }
+}