3 files changed, 79 insertions, 16 deletions

diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
index dea552a60..a841b97bf 100644
--- a/crates/hir_def/src/attr.rs
+++ b/crates/hir_def/src/attr.rs
@@ -171,6 +171,9 @@ pub struct AttrQuery<'a> {
 }
 
 impl<'a> AttrQuery<'a> {
+    /// For an attribute like `#[attr(value)]`, returns the `(value)` subtree.
+    ///
+    /// If the attribute does not have a token tree argument, returns `None`.
     pub fn tt_values(self) -> impl Iterator<Item = &'a Subtree> {
         self.attrs().filter_map(|attr| match attr.input.as_ref()? {
             AttrInput::TokenTree(it) => Some(it),
@@ -178,6 +181,9 @@ impl<'a> AttrQuery<'a> {
         })
     }
 
+    /// For an attribute like `#[key = "value"]`, returns `"value"`.
+    ///
+    /// Returns `None` if the attribute does not have `key = "value"` form.
     pub fn string_value(self) -> Option<&'a SmolStr> {
         self.attrs().find_map(|attr| match attr.input.as_ref()? {
             AttrInput::Literal(it) => Some(it),
diff --git a/crates/hir_def/src/nameres/collector.rs b/crates/hir_def/src/nameres/collector.rs
index 4c3993ff0..42c0f0536 100644
--- a/crates/hir_def/src/nameres/collector.rs
+++ b/crates/hir_def/src/nameres/collector.rs
@@ -16,10 +16,10 @@ use hir_expand::{
     proc_macro::ProcMacroExpander,
     HirFileId, MacroCallId, MacroDefId, MacroDefKind,
 };
-use rustc_hash::FxHashMap;
-use rustc_hash::FxHashSet;
+use rustc_hash::{FxHashMap, FxHashSet};
 use syntax::ast;
 use test_utils::mark;
+use tt::{Leaf, TokenTree};
 
 use crate::{
     attr::Attrs,
@@ -281,6 +281,25 @@ impl DefCollector<'_> {
         }
     }
 
+    fn resolve_proc_macro(&mut self, name: &Name) {
+        let macro_def = match self.proc_macros.iter().find(|(n, _)| n == name) {
+            Some((_, expander)) => MacroDefId {
+                ast_id: None,
+                krate: Some(self.def_map.krate),
+                kind: MacroDefKind::ProcMacro(*expander),
+                local_inner: false,
+            },
+            None => MacroDefId {
+                ast_id: None,
+                krate: Some(self.def_map.krate),
+                kind: MacroDefKind::ProcMacro(ProcMacroExpander::dummy(self.def_map.krate)),
+                local_inner: false,
+            },
+        };
+
+        self.define_proc_macro(name.clone(), macro_def);
+    }
+
     /// Define a macro with `macro_rules`.
     ///
     /// It will define the macro in legacy textual scope, and if it has `#[macro_export]`,
@@ -917,6 +936,9 @@ impl ModCollector<'_, '_> {
                 }
                 ModItem::Function(id) => {
                     let func = &self.item_tree[id];
+
+                    self.collect_proc_macro_def(&func.name, attrs);
+
                     def = Some(DefData {
                         id: FunctionLoc {
                             container: container.into(),
@@ -1177,6 +1199,30 @@ impl ModCollector<'_, '_> {
         }
     }
 
+    /// If `attrs` registers a procedural macro, collects its definition.
+    fn collect_proc_macro_def(&mut self, func_name: &Name, 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 {
+                    return;
+                }
+            } else {
+                return;
+            }
+        };
+
+        self.def_collector.resolve_proc_macro(&macro_name);
+    }
+
     fn collect_macro(&mut self, mac: &MacroCall) {
         let mut ast_id = AstIdWithPath::new(self.file_id, mac.ast_id, mac.path.clone());
 
diff --git a/crates/hir_expand/src/proc_macro.rs b/crates/hir_expand/src/proc_macro.rs
index 80255ea32..7505cb061 100644
--- a/crates/hir_expand/src/proc_macro.rs
+++ b/crates/hir_expand/src/proc_macro.rs
@@ -7,7 +7,7 @@ use tt::buffer::{Cursor, TokenBuffer};
 #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
 pub struct ProcMacroExpander {
     krate: CrateId,
-    proc_macro_id: ProcMacroId,
+    proc_macro_id: Option<ProcMacroId>,
 }
 
 macro_rules! err {
@@ -20,8 +20,14 @@ macro_rules! err {
 }
 
 impl ProcMacroExpander {
-    pub fn new(krate: CrateId, proc_macro_id: ProcMacroId) -> ProcMacroExpander {
-        ProcMacroExpander { krate, proc_macro_id }
+    pub fn new(krate: CrateId, proc_macro_id: ProcMacroId) -> Self {
+        Self { krate, proc_macro_id: Some(proc_macro_id) }
+    }
+
+    pub fn dummy(krate: CrateId) -> Self {
+        // FIXME: Should store the name for better errors
+        // FIXME: I think this is the second layer of "dummy" expansion, we should reduce that
+        Self { krate, proc_macro_id: None }
     }
 
     pub fn expand(
@@ -30,17 +36,22 @@ impl ProcMacroExpander {
         _id: LazyMacroId,
         tt: &tt::Subtree,
     ) -> Result<tt::Subtree, mbe::ExpandError> {
-        let krate_graph = db.crate_graph();
-        let proc_macro = krate_graph[self.krate]
-            .proc_macro
-            .get(self.proc_macro_id.0 as usize)
-            .clone()
-            .ok_or_else(|| err!("No derive macro found."))?;
-
-        let tt = remove_derive_attrs(tt)
-            .ok_or_else(|| err!("Fail to remove derive for custom derive"))?;
-
-        proc_macro.expander.expand(&tt, None).map_err(mbe::ExpandError::from)
+        match self.proc_macro_id {
+            Some(id) => {
+                let krate_graph = db.crate_graph();
+                let proc_macro = krate_graph[self.krate]
+                    .proc_macro
+                    .get(id.0 as usize)
+                    .clone()
+                    .ok_or_else(|| err!("No derive macro found."))?;
+
+                let tt = remove_derive_attrs(tt)
+                    .ok_or_else(|| err!("Fail to remove derive for custom derive"))?;
+
+                proc_macro.expander.expand(&tt, None).map_err(mbe::ExpandError::from)
+            }
+            None => Err(err!("Unresolved proc macro")),
+        }
     }
 }