Use hir_def to resolve proc macros

1 files changed, 48 insertions, 2 deletions

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());