8 files changed, 183 insertions, 61 deletions

diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
index e4c84afbf..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),
@@ -283,8 +283,51 @@ impl Attrs {
     /// Constructs a map that maps the lowered `Attr`s in this `Attrs` back to its original syntax nodes.
     ///
     /// `owner` must be the original owner of the attributes.
-    pub fn source_map(&self, owner: &dyn ast::AttrsOwner) -> AttrSourceMap {
-        AttrSourceMap { attrs: collect_attrs(owner).collect() }
+    // FIXME: figure out a better api that doesnt require the for_module hack
+    pub fn source_map(&self, owner: InFile<&dyn ast::AttrsOwner>) -> AttrSourceMap {
+        // FIXME: This doesn't work correctly for modules, as the attributes there can have up to
+        // two different owners
+        AttrSourceMap {
+            attrs: collect_attrs(owner.value)
+                .map(|attr| InFile::new(owner.file_id, attr))
+                .collect(),
+        }
+    }
+
+    pub fn source_map_for_module(
+        &self,
+        db: &dyn DefDatabase,
+        module: crate::ModuleId,
+    ) -> AttrSourceMap {
+        let def_map = module.def_map(db);
+        let mod_data = &def_map[module.local_id];
+        let attrs = match mod_data.declaration_source(db) {
+            Some(it) => {
+                let mut attrs: Vec<_> = collect_attrs(&it.value as &dyn ast::AttrsOwner)
+                    .map(|attr| InFile::new(it.file_id, attr))
+                    .collect();
+                if let InFile { file_id, value: ModuleSource::SourceFile(file) } =
+                    mod_data.definition_source(db)
+                {
+                    attrs.extend(
+                        collect_attrs(&file as &dyn ast::AttrsOwner)
+                            .map(|attr| InFile::new(file_id, attr)),
+                    )
+                }
+                attrs
+            }
+            None => {
+                let InFile { file_id, value } = mod_data.definition_source(db);
+                match &value {
+                    ModuleSource::SourceFile(file) => collect_attrs(file as &dyn ast::AttrsOwner),
+                    ModuleSource::Module(module) => collect_attrs(module as &dyn ast::AttrsOwner),
+                    ModuleSource::BlockExpr(block) => collect_attrs(block as &dyn ast::AttrsOwner),
+                }
+                .map(|attr| InFile::new(file_id, attr))
+                .collect()
+            }
+        };
+        AttrSourceMap { attrs }
     }
 
     pub fn by_key(&self, key: &'static str) -> AttrQuery<'_> {
@@ -379,7 +422,7 @@ fn inner_attributes(
 }
 
 pub struct AttrSourceMap {
-    attrs: Vec<Either<ast::Attr, ast::Comment>>,
+    attrs: Vec<InFile<Either<ast::Attr, ast::Comment>>>,
 }
 
 impl AttrSourceMap {
@@ -389,10 +432,11 @@ impl AttrSourceMap {
     ///
     /// Note that the returned syntax node might be a `#[cfg_attr]`, or a doc comment, instead of
     /// the attribute represented by `Attr`.
-    pub fn source_of(&self, attr: &Attr) -> &Either<ast::Attr, ast::Comment> {
+    pub fn source_of(&self, attr: &Attr) -> InFile<&Either<ast::Attr, ast::Comment>> {
         self.attrs
             .get(attr.index as usize)
             .unwrap_or_else(|| panic!("cannot find `Attr` at index {}", attr.index))
+            .as_ref()
     }
 }
 
@@ -428,18 +472,6 @@ impl Attr {
         Some(Attr { index, path, input })
     }
 
-    /// Maps this lowered `Attr` back to its original syntax node.
-    ///
-    /// `owner` must be the original owner of the attribute.
-    ///
-    /// Note that the returned syntax node might be a `#[cfg_attr]`, or a doc comment, instead of
-    /// the attribute represented by `Attr`.
-    pub fn to_src(&self, owner: &dyn ast::AttrsOwner) -> Either<ast::Attr, ast::Comment> {
-        collect_attrs(owner).nth(self.index as usize).unwrap_or_else(|| {
-            panic!("cannot find `Attr` at index {} in {}", self.index, owner.syntax())
-        })
-    }
-
     /// Parses this attribute as a `#[derive]`, returns an iterator that yields all contained paths
     /// to derive macros.
     ///
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/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!(),
+    }
+}