10 files changed, 147 insertions, 45 deletions

diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
index 739e3f5e3..1cab0e363 100644
--- a/crates/hir_def/src/attr.rs
+++ b/crates/hir_def/src/attr.rs
@@ -209,7 +209,7 @@ impl Attrs {
             },
             AttrDefId::TraitId(it) => attrs_from_item_tree(it.lookup(db).id, db),
             AttrDefId::MacroDefId(it) => {
-                it.ast_id.map_or_else(Default::default, |ast_id| attrs_from_ast(ast_id, db))
+                it.ast_id().map_or_else(Default::default, |ast_id| attrs_from_ast(ast_id, db))
             }
             AttrDefId::ImplId(it) => attrs_from_item_tree(it.lookup(db).id, db),
             AttrDefId::ConstId(it) => attrs_from_item_tree(it.lookup(db).id, db),
diff --git a/crates/hir_def/src/import_map.rs b/crates/hir_def/src/import_map.rs
index 369bc3350..960cabb5f 100644
--- a/crates/hir_def/src/import_map.rs
+++ b/crates/hir_def/src/import_map.rs
@@ -912,10 +912,10 @@ mod tests {
                 dep::fmt (t)
                 dep::format (f)
                 dep::Fmt (v)
-                dep::fmt::Display (t)
+                dep::Fmt (m)
                 dep::Fmt (t)
                 dep::fmt::Display::fmt (a)
-                dep::Fmt (m)
+                dep::fmt::Display (t)
             "#]],
         );
 
@@ -926,9 +926,9 @@ mod tests {
             expect![[r#"
                 dep::fmt (t)
                 dep::Fmt (v)
+                dep::Fmt (m)
                 dep::Fmt (t)
                 dep::fmt::Display::fmt (a)
-                dep::Fmt (m)
             "#]],
         );
 
@@ -939,10 +939,10 @@ mod tests {
             expect![[r#"
                 dep::fmt (t)
                 dep::Fmt (v)
-                dep::fmt::Display (t)
+                dep::Fmt (m)
                 dep::Fmt (t)
                 dep::fmt::Display::fmt (a)
-                dep::Fmt (m)
+                dep::fmt::Display (t)
             "#]],
         );
     }
@@ -980,10 +980,10 @@ mod tests {
             expect![[r#"
                 dep::fmt (t)
                 dep::Fmt (v)
-                dep::fmt::Display (t)
+                dep::Fmt (m)
                 dep::Fmt (t)
                 dep::fmt::Display::fmt (a)
-                dep::Fmt (m)
+                dep::fmt::Display (t)
             "#]],
         );
 
@@ -994,9 +994,9 @@ mod tests {
             expect![[r#"
                 dep::fmt (t)
                 dep::Fmt (v)
+                dep::Fmt (m)
                 dep::Fmt (t)
                 dep::fmt::Display::fmt (a)
-                dep::Fmt (m)
             "#]],
         );
     }
@@ -1058,8 +1058,8 @@ mod tests {
             Query::new("".to_string()).limit(2),
             expect![[r#"
                 dep::fmt (t)
-                dep::Fmt (t)
                 dep::Fmt (m)
+                dep::Fmt (t)
                 dep::Fmt (v)
             "#]],
         );
diff --git a/crates/hir_def/src/item_scope.rs b/crates/hir_def/src/item_scope.rs
index aafd73b60..f3ebe7c72 100644
--- a/crates/hir_def/src/item_scope.rs
+++ b/crates/hir_def/src/item_scope.rs
@@ -252,7 +252,7 @@ impl ItemScope {
             .for_each(|vis| *vis = Visibility::Module(this_module));
 
         for (mac, vis) in self.macros.values_mut() {
-            if let MacroDefKind::ProcMacro(_) = mac.kind {
+            if let MacroDefKind::ProcMacro(..) = mac.kind {
                 // FIXME: Technically this is insufficient since reexports of proc macros are also
                 // forbidden. Practically nobody does that.
                 continue;
diff --git a/crates/hir_def/src/item_tree.rs b/crates/hir_def/src/item_tree.rs
index 5f5b7151a..ae2475b4e 100644
--- a/crates/hir_def/src/item_tree.rs
+++ b/crates/hir_def/src/item_tree.rs
@@ -115,6 +115,10 @@ impl ItemTree {
                     // still need to collect inner items.
                     ctx.lower_inner_items(stmt.syntax())
                 },
+                ast::Item(item) => {
+                    // Macros can expand to stmt and other item, and we add it as top level item
+                    ctx.lower_single_item(item)
+                },
                 _ => {
                     panic!("cannot create item tree from {:?} {}", syntax, syntax);
                 },
diff --git a/crates/hir_def/src/item_tree/lower.rs b/crates/hir_def/src/item_tree/lower.rs
index 3f558edd8..d3fe1ce1e 100644
--- a/crates/hir_def/src/item_tree/lower.rs
+++ b/crates/hir_def/src/item_tree/lower.rs
@@ -87,6 +87,14 @@ impl Ctx {
         self.tree
     }
 
+    pub(super) fn lower_single_item(mut self, item: ast::Item) -> ItemTree {
+        self.tree.top_level = self
+            .lower_mod_item(&item, false)
+            .map(|item| item.0)
+            .unwrap_or_else(|| Default::default());
+        self.tree
+    }
+
     pub(super) fn lower_inner_items(mut self, within: &SyntaxNode) -> ItemTree {
         self.collect_inner_items(within);
         self.tree
diff --git a/crates/hir_def/src/lib.rs b/crates/hir_def/src/lib.rs
index 6758411a0..21add086d 100644
--- a/crates/hir_def/src/lib.rs
+++ b/crates/hir_def/src/lib.rs
@@ -650,7 +650,7 @@ fn macro_call_as_call_id(
 ) -> Result<Result<MacroCallId, ErrorEmitted>, UnresolvedMacro> {
     let def: MacroDefId = resolver(call.path.clone()).ok_or(UnresolvedMacro)?;
 
-    let res = if let MacroDefKind::BuiltInEager(_) = def.kind {
+    let res = if let MacroDefKind::BuiltInEager(..) = def.kind {
         let macro_call = InFile::new(call.ast_id.file_id, call.ast_id.to_node(db.upcast()));
         let hygiene = Hygiene::new(db.upcast(), call.ast_id.file_id);
 
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 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!(),
+    }
+}