Record custom derive helpers in `DefMap`

Also clean up proc macro attribute parsing a bit

3 files changed, 91 insertions, 25 deletions

diff --git a/crates/hir_def/src/nameres.rs b/crates/hir_def/src/nameres.rs
index c97be584e..1ac326f97 100644
--- a/crates/hir_def/src/nameres.rs
+++ b/crates/hir_def/src/nameres.rs
@@ -53,11 +53,12 @@ mod path_resolution;
 
 #[cfg(test)]
 mod tests;
+mod proc_macro;
 
 use std::sync::Arc;
 
 use base_db::{CrateId, Edition, FileId};
-use hir_expand::{diagnostics::DiagnosticSink, name::Name, InFile};
+use hir_expand::{diagnostics::DiagnosticSink, name::Name, InFile, MacroDefId};
 use la_arena::Arena;
 use profile::Count;
 use rustc_hash::FxHashMap;
@@ -73,6 +74,8 @@ use crate::{
     AstId, BlockId, BlockLoc, LocalModuleId, ModuleDefId, ModuleId,
 };
 
+use self::proc_macro::ProcMacroDef;
+
 /// Contains the results of (early) name resolution.
 ///
 /// A `DefMap` stores the module tree and the definitions that are in scope in every module after
@@ -95,6 +98,12 @@ pub struct DefMap {
     prelude: Option<ModuleId>,
     extern_prelude: FxHashMap<Name, ModuleDefId>,
 
+    /// Side table with additional proc. macro info, for use by name resolution in downstream
+    /// crates.
+    ///
+    /// (the primary purpose is to resolve derive helpers)
+    exported_proc_macros: FxHashMap<MacroDefId, ProcMacroDef>,
+
     edition: Edition,
     diagnostics: Vec<DefDiagnostic>,
 }
@@ -237,6 +246,7 @@ impl DefMap {
             krate,
             edition,
             extern_prelude: FxHashMap::default(),
+            exported_proc_macros: FxHashMap::default(),
             prelude: None,
             root,
             modules,
diff --git a/crates/hir_def/src/nameres/collector.rs b/crates/hir_def/src/nameres/collector.rs
index 696370ada..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,9 +354,9 @@ 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, ast_id: AstId<ast::Fn>) {
+    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 {
                 krate: self.def_map.krate,
                 kind: MacroDefKind::ProcMacro(*expander, ast_id),
@@ -368,7 +369,8 @@ impl DefCollector<'_> {
             },
         };
 
-        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`.
@@ -1386,26 +1388,9 @@ impl ModCollector<'_, '_> {
     /// If `attrs` registers a procedural macro, collects its definition.
     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, ast_id);
+        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>) {
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
+        }
+    }
+}