Merge pull request #19862 from ChayimFriedman2/item-resolve-macro-hir

fix: Fix IDE resolution of item macros

Author: ChayimFriedman2

crates/hir-def/src/db.rs

@@ -422,6 +422,7 @@ fn macro_def(db: &dyn DefDatabase, id: MacroId) -> MacroDefId {
             let makro = &item_tree[loc.id.value];
             MacroDefId {
                 krate: loc.container.krate,
+                block: loc.container.block.map(|block| salsa::plumbing::AsId::as_id(&block)),
                 kind: kind(loc.expander, loc.id.file_id(), makro.ast_id.upcast()),
                 local_inner: false,
                 allow_internal_unsafe: loc.allow_internal_unsafe,
@@ -435,6 +436,7 @@ fn macro_def(db: &dyn DefDatabase, id: MacroId) -> MacroDefId {
             let makro = &item_tree[loc.id.value];
             MacroDefId {
                 krate: loc.container.krate,
+                block: loc.container.block.map(|block| salsa::plumbing::AsId::as_id(&block)),
                 kind: kind(loc.expander, loc.id.file_id(), makro.ast_id.upcast()),
                 local_inner: loc.flags.contains(MacroRulesLocFlags::LOCAL_INNER),
                 allow_internal_unsafe: loc
@@ -450,6 +452,7 @@ fn macro_def(db: &dyn DefDatabase, id: MacroId) -> MacroDefId {
             let makro = &item_tree[loc.id.value];
             MacroDefId {
                 krate: loc.container.krate,
+                block: None,
                 kind: MacroDefKind::ProcMacro(
                     InFile::new(loc.id.file_id(), makro.ast_id),
                     loc.expander,

crates/hir-def/src/resolver.rs

@@ -696,6 +696,15 @@ impl<'db> Resolver<'db> {
         &def_map[local_id].scope
     }
 
+    pub fn item_scopes(&self) -> impl Iterator<Item = &ItemScope> {
+        self.scopes()
+            .filter_map(move |scope| match scope {
+                Scope::BlockScope(m) => Some(&m.def_map[m.module_id].scope),
+                _ => None,
+            })
+            .chain(std::iter::once(&self.module_scope.def_map[self.module_scope.module_id].scope))
+    }
+
     pub fn krate(&self) -> Crate {
         self.module_scope.def_map.krate()
     }

crates/hir-expand/src/lib.rs

@@ -258,6 +258,8 @@ pub struct MacroCallLoc {
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct MacroDefId {
     pub krate: Crate,
+    // FIXME: In `hir-expand` we can't refer to `BlockId`.
+    pub block: Option<salsa::Id>,
     pub edition: Edition,
     pub kind: MacroDefKind,
     pub local_inner: bool,

crates/hir/src/semantics.rs

@@ -411,7 +411,7 @@ impl<'db> SemanticsImpl<'db> {
         let sa = self.analyze_no_infer(macro_call.syntax())?;
 
         let macro_call = InFile::new(sa.file_id, macro_call);
-        let file_id = sa.expand(self.db, macro_call)?;
+        let file_id = sa.expansion(self.db, macro_call)?;
 
         let node = self.parse_or_expand(file_id.into());
         Some(InFile::new(file_id.into(), node))
@@ -437,7 +437,7 @@ impl<'db> SemanticsImpl<'db> {
         let sa = self.analyze_no_infer(macro_call.syntax())?;
 
         let macro_call = InFile::new(sa.file_id, macro_call);
-        let file_id = sa.expand(self.db, macro_call)?;
+        let file_id = sa.expansion(self.db, macro_call)?;
         let macro_call = self.db.lookup_intern_macro_call(file_id);
 
         let skip = matches!(
@@ -576,7 +576,7 @@ impl<'db> SemanticsImpl<'db> {
     ) -> Option<(SyntaxNode, Vec<(SyntaxToken, u8)>)> {
         let analyzer = self.analyze_no_infer(actual_macro_call.syntax())?;
         let macro_call = InFile::new(analyzer.file_id, actual_macro_call);
-        let macro_file = analyzer.expansion(macro_call)?;
+        let macro_file = analyzer.expansion(self.db, macro_call)?;
         hir_expand::db::expand_speculative(
             self.db,
             macro_file,
@@ -1106,7 +1106,7 @@ impl<'db> SemanticsImpl<'db> {
                                                 false,
                                             )
                                         })?
-                                        .expand(self.db, mcall.as_ref())?;
+                                        .expansion(self.db, mcall.as_ref())?;
                                     m_cache.insert(mcall, it);
                                     it
                                 }
@@ -1565,7 +1565,7 @@ impl<'db> SemanticsImpl<'db> {
         let sa = self.analyze(macro_call.syntax())?;
         self.db
             .parse_macro_expansion(
-                sa.expand(self.db, self.wrap_node_infile(macro_call.clone()).as_ref())?,
+                sa.expansion(self.db, self.wrap_node_infile(macro_call.clone()).as_ref())?,
             )
             .value
             .1

crates/hir/src/source_analyzer.rs

@@ -26,7 +26,7 @@ use hir_def::{
     },
     hir::{BindingId, Expr, ExprId, ExprOrPatId, Pat},
     lang_item::LangItem,
-    nameres::{MacroSubNs, crate_def_map},
+    nameres::{MacroSubNs, block_def_map, crate_def_map},
     resolver::{HasResolver, Resolver, TypeNs, ValueNs, resolver_for_scope},
     type_ref::{Mutability, TypeRefId},
 };
@@ -218,8 +218,16 @@ impl<'db> SourceAnalyzer<'db> {
         })
     }
 
-    pub(crate) fn expansion(&self, node: InFile<&ast::MacroCall>) -> Option<MacroCallId> {
-        self.store_sm()?.expansion(node)
+    pub(crate) fn expansion(
+        &self,
+        db: &dyn HirDatabase,
+        macro_call: InFile<&ast::MacroCall>,
+    ) -> Option<MacroCallId> {
+        self.store_sm().and_then(|sm| sm.expansion(macro_call)).or_else(|| {
+            let ast_id_map = db.ast_id_map(macro_call.file_id);
+            let call_ast_id = macro_call.with_value(ast_id_map.ast_id(macro_call.value));
+            self.resolver.item_scopes().find_map(|scope| scope.macro_invoc(call_ast_id))
+        })
     }
 
     fn trait_environment(&self, db: &'db dyn HirDatabase) -> Arc<TraitEnvironment> {
@@ -747,17 +755,16 @@ impl<'db> SourceAnalyzer<'db> {
 
     pub(crate) fn resolve_macro_call(
         &self,
-        db: &'db dyn HirDatabase,
+        db: &dyn HirDatabase,
         macro_call: InFile<&ast::MacroCall>,
     ) -> Option<Macro> {
-        let bs = self.store_sm()?;
-        bs.expansion(macro_call).and_then(|it| {
-            // FIXME: Block def maps
+        self.expansion(db, macro_call).and_then(|it| {
             let def = it.lookup(db).def;
-            crate_def_map(db, def.krate)
-                .macro_def_to_macro_id
-                .get(&def.kind.erased_ast_id())
-                .map(|it| (*it).into())
+            let def_map = match def.block {
+                Some(block) => block_def_map(db, base_db::salsa::plumbing::FromId::from_id(block)),
+                None => crate_def_map(db, def.krate),
+            };
+            def_map.macro_def_to_macro_id.get(&def.kind.erased_ast_id()).map(|it| (*it).into())
         })
     }
 
@@ -1292,18 +1299,6 @@ impl<'db> SourceAnalyzer<'db> {
             .collect()
     }
 
-    pub(crate) fn expand(
-        &self,
-        db: &'db dyn HirDatabase,
-        macro_call: InFile<&ast::MacroCall>,
-    ) -> Option<MacroCallId> {
-        self.store_sm().and_then(|bs| bs.expansion(macro_call)).or_else(|| {
-            self.resolver.item_scope().macro_invoc(
-                macro_call.with_value(db.ast_id_map(macro_call.file_id).ast_id(macro_call.value)),
-            )
-        })
-    }
-
     pub(crate) fn resolve_variant(&self, record_lit: ast::RecordExpr) -> Option<VariantId> {
         let infer = self.infer()?;
         let expr_id = self.expr_id(record_lit.into())?;

crates/ide-completion/src/tests/item.rs

@@ -4,7 +4,7 @@
 //! in [crate::completions::mod_].
 use expect_test::expect;
 
-use crate::tests::{check_edit, check_with_base_items};
+use crate::tests::{check, check_edit, check_with_base_items};
 
 #[test]
 fn target_type_or_trait_in_impl_block() {
@@ -308,3 +308,39 @@ fn bar() {
         "#]],
     );
 }
+
+#[test]
+fn expression_in_item_macro() {
+    check(
+        r#"
+fn foo() -> u8 { 0 }
+
+macro_rules! foo {
+    ($expr:expr) => {
+        const BAR: u8 = $expr;
+    };
+}
+
+foo!(f$0);
+    "#,
+        expect![[r#"
+            ct BAR                   u8
+            fn foo()         fn() -> u8
+            ma foo!(…) macro_rules! foo
+            bt u32                  u32
+            kw const
+            kw crate::
+            kw false
+            kw for
+            kw if
+            kw if let
+            kw loop
+            kw match
+            kw self::
+            kw true
+            kw unsafe
+            kw while
+            kw while let
+        "#]],
+    );
+}