Auto merge of #16679 - Veykril:block-mod, r=Veykril

fix: Fix modules in blocks not resolving in ide layer

Fixes https://github.com/rust-lang/rust-analyzer/issues/16511

11 files changed, 183 insertions, 33 deletions

diff --git a/crates/base-db/src/input.rs b/crates/base-db/src/input.rs
index a817cd0c3ac..b243b37b77b 100644
--- a/crates/base-db/src/input.rs
+++ b/crates/base-db/src/input.rs
@@ -570,7 +570,7 @@ impl CrateGraph {
                 .arena
                 .iter_mut()
                 .take(m)
-                .find_map(|(id, data)| merge((id, data), (topo, &crate_data)).then_some(id));
+                .find_map(|(id, data)| merge((id, data), (topo, crate_data)).then_some(id));
 
             let new_id =
                 if let Some(res) = res { res } else { self.arena.alloc(crate_data.clone()) };
diff --git a/crates/hir-def/src/body/tests/block.rs b/crates/hir-def/src/body/tests/block.rs
index 44eeed9e3fb..985c6387ba0 100644
--- a/crates/hir-def/src/body/tests/block.rs
+++ b/crates/hir-def/src/body/tests/block.rs
@@ -299,6 +299,40 @@ pub mod cov_mark {
 }
 
 #[test]
+fn macro_exported_in_block_mod() {
+    check_at(
+        r#"
+#[macro_export]
+macro_rules! foo {
+    () => { pub struct FooWorks; };
+}
+macro_rules! bar {
+    () => { pub struct BarWorks; };
+}
+fn main() {
+    mod module {
+        foo!();
+        bar!();
+        $0
+    }
+}
+"#,
+        expect![[r#"
+            block scope
+            module: t
+
+            block scope::module
+            BarWorks: t v
+            FooWorks: t v
+
+            crate
+            foo: m
+            main: v
+        "#]],
+    );
+}
+
+#[test]
 fn macro_resolve_legacy() {
     check_at(
         r#"
diff --git a/crates/hir-def/src/child_by_source.rs b/crates/hir-def/src/child_by_source.rs
index ba7d06272af..f1c6b3b89fc 100644
--- a/crates/hir-def/src/child_by_source.rs
+++ b/crates/hir-def/src/child_by_source.rs
@@ -189,10 +189,11 @@ impl ChildBySource for DefWithBodyId {
             VariantId::EnumVariantId(v).child_by_source_to(db, res, file_id)
         }
 
-        for (_, def_map) in body.blocks(db) {
+        for (block, def_map) in body.blocks(db) {
             // All block expressions are merged into the same map, because they logically all add
             // inner items to the containing `DefWithBodyId`.
             def_map[DefMap::ROOT].scope.child_by_source_to(db, res, file_id);
+            res[keys::BLOCK].insert(block.lookup(db).ast_id.to_node(db.upcast()), block);
         }
     }
 }
diff --git a/crates/hir-def/src/dyn_map/keys.rs b/crates/hir-def/src/dyn_map/keys.rs
index 60832f59eb9..f83ab1e1a05 100644
--- a/crates/hir-def/src/dyn_map/keys.rs
+++ b/crates/hir-def/src/dyn_map/keys.rs
@@ -8,13 +8,14 @@ use syntax::{ast, AstNode, AstPtr};
 
 use crate::{
     dyn_map::{DynMap, Policy},
-    ConstId, EnumId, EnumVariantId, ExternCrateId, FieldId, FunctionId, ImplId, LifetimeParamId,
-    Macro2Id, MacroRulesId, ProcMacroId, StaticId, StructId, TraitAliasId, TraitId, TypeAliasId,
-    TypeOrConstParamId, UnionId, UseId,
+    BlockId, ConstId, EnumId, EnumVariantId, ExternCrateId, FieldId, FunctionId, ImplId,
+    LifetimeParamId, Macro2Id, MacroRulesId, ProcMacroId, StaticId, StructId, TraitAliasId,
+    TraitId, TypeAliasId, TypeOrConstParamId, UnionId, UseId,
 };
 
 pub type Key<K, V> = crate::dyn_map::Key<K, V, AstPtrPolicy<K, V>>;
 
+pub const BLOCK: Key<ast::BlockExpr, BlockId> = Key::new();
 pub const FUNCTION: Key<ast::Fn, FunctionId> = Key::new();
 pub const CONST: Key<ast::Const, ConstId> = Key::new();
 pub const STATIC: Key<ast::Static, StaticId> = Key::new();
diff --git a/crates/hir-def/src/nameres.rs b/crates/hir-def/src/nameres.rs
index a2eca066438..1a827821156 100644
--- a/crates/hir-def/src/nameres.rs
+++ b/crates/hir-def/src/nameres.rs
@@ -469,6 +469,12 @@ impl DefMap {
         CrateRootModuleId { krate: self.krate }
     }
 
+    /// This is the same as [`Self::crate_root`] for crate def maps, but for block def maps, it
+    /// returns the root block module.
+    pub fn root_module_id(&self) -> ModuleId {
+        self.module_id(Self::ROOT)
+    }
+
     pub(crate) fn resolve_path(
         &self,
         db: &dyn DefDatabase,
diff --git a/crates/hir-def/src/resolver.rs b/crates/hir-def/src/resolver.rs
index db47d743c5a..226d6f513f5 100644
--- a/crates/hir-def/src/resolver.rs
+++ b/crates/hir-def/src/resolver.rs
@@ -1,5 +1,5 @@
 //! Name resolution façade.
-use std::{fmt, hash::BuildHasherDefault};
+use std::{fmt, hash::BuildHasherDefault, mem};
 
 use base_db::CrateId;
 use hir_expand::{
@@ -809,7 +809,7 @@ fn resolver_for_scope_(
     for scope in scope_chain.into_iter().rev() {
         if let Some(block) = scopes.block(scope) {
             let def_map = db.block_def_map(block);
-            r = r.push_block_scope(def_map, DefMap::ROOT);
+            r = r.push_block_scope(def_map);
             // FIXME: This adds as many module scopes as there are blocks, but resolving in each
             // already traverses all parents, so this is O(n²). I think we could only store the
             // innermost module scope instead?
@@ -835,8 +835,9 @@ impl Resolver {
         self.push_scope(Scope::ImplDefScope(impl_def))
     }
 
-    fn push_block_scope(self, def_map: Arc<DefMap>, module_id: LocalModuleId) -> Resolver {
-        self.push_scope(Scope::BlockScope(ModuleItemMap { def_map, module_id }))
+    fn push_block_scope(self, def_map: Arc<DefMap>) -> Resolver {
+        debug_assert!(def_map.block_id().is_some());
+        self.push_scope(Scope::BlockScope(ModuleItemMap { def_map, module_id: DefMap::ROOT }))
     }
 
     fn push_expr_scope(
@@ -986,19 +987,27 @@ pub trait HasResolver: Copy {
 impl HasResolver for ModuleId {
     fn resolver(self, db: &dyn DefDatabase) -> Resolver {
         let mut def_map = self.def_map(db);
-        let mut modules: SmallVec<[_; 1]> = smallvec![];
         let mut module_id = self.local_id;
+        let mut modules: SmallVec<[_; 1]> = smallvec![];
+
+        if !self.is_block_module() {
+            return Resolver { scopes: vec![], module_scope: ModuleItemMap { def_map, module_id } };
+        }
+
         while let Some(parent) = def_map.parent() {
-            modules.push((def_map, module_id));
-            def_map = parent.def_map(db);
-            module_id = parent.local_id;
+            let block_def_map = mem::replace(&mut def_map, parent.def_map(db));
+            modules.push(block_def_map);
+            if !parent.is_block_module() {
+                module_id = parent.local_id;
+                break;
+            }
         }
         let mut resolver = Resolver {
             scopes: Vec::with_capacity(modules.len()),
             module_scope: ModuleItemMap { def_map, module_id },
         };
-        for (def_map, module) in modules.into_iter().rev() {
-            resolver = resolver.push_block_scope(def_map, module);
+        for def_map in modules.into_iter().rev() {
+            resolver = resolver.push_block_scope(def_map);
         }
         resolver
     }
diff --git a/crates/hir/src/semantics/source_to_def.rs b/crates/hir/src/semantics/source_to_def.rs
index 14dbe692403..ef4ed90ce35 100644
--- a/crates/hir/src/semantics/source_to_def.rs
+++ b/crates/hir/src/semantics/source_to_def.rs
@@ -86,6 +86,7 @@
 //! syntax nodes against this specific crate.
 
 use base_db::FileId;
+use either::Either;
 use hir_def::{
     child_by_source::ChildBySource,
     dyn_map::{
@@ -93,9 +94,9 @@ use hir_def::{
         DynMap,
     },
     hir::{BindingId, LabelId},
-    AdtId, ConstId, ConstParamId, DefWithBodyId, EnumId, EnumVariantId, ExternCrateId, FieldId,
-    FunctionId, GenericDefId, GenericParamId, ImplId, LifetimeParamId, MacroId, ModuleId, StaticId,
-    StructId, TraitAliasId, TraitId, TypeAliasId, TypeParamId, UnionId, UseId, VariantId,
+    AdtId, BlockId, ConstId, ConstParamId, DefWithBodyId, EnumId, EnumVariantId, ExternCrateId,
+    FieldId, FunctionId, GenericDefId, GenericParamId, ImplId, LifetimeParamId, MacroId, ModuleId,
+    StaticId, StructId, TraitAliasId, TraitId, TypeAliasId, TypeParamId, UnionId, UseId, VariantId,
 };
 use hir_expand::{attrs::AttrId, name::AsName, HirFileId, HirFileIdExt, MacroCallId};
 use rustc_hash::FxHashMap;
@@ -131,15 +132,19 @@ impl SourceToDefCtx<'_, '_> {
         mods
     }
 
-    pub(super) fn module_to_def(&self, src: InFile<ast::Module>) -> Option<ModuleId> {
+    pub(super) fn module_to_def(&mut self, src: InFile<ast::Module>) -> Option<ModuleId> {
         let _p = tracing::span!(tracing::Level::INFO, "module_to_def");
         let parent_declaration = src
             .syntax()
             .ancestors_with_macros_skip_attr_item(self.db.upcast())
-            .find_map(|it| it.map(ast::Module::cast).transpose());
+            .find_map(|it| it.map(Either::<ast::Module, ast::BlockExpr>::cast).transpose())
+            .map(|it| it.transpose());
 
         let parent_module = match parent_declaration {
-            Some(parent_declaration) => self.module_to_def(parent_declaration),
+            Some(Either::Right(parent_block)) => self
+                .block_to_def(parent_block)
+                .map(|block| self.db.block_def_map(block).root_module_id()),
+            Some(Either::Left(parent_declaration)) => self.module_to_def(parent_declaration),
             None => {
                 let file_id = src.file_id.original_file(self.db.upcast());
                 self.file_to_def(file_id).first().copied()
@@ -197,6 +202,9 @@ impl SourceToDefCtx<'_, '_> {
     pub(super) fn tuple_field_to_def(&mut self, src: InFile<ast::TupleField>) -> Option<FieldId> {
         self.to_def(src, keys::TUPLE_FIELD)
     }
+    pub(super) fn block_to_def(&mut self, src: InFile<ast::BlockExpr>) -> Option<BlockId> {
+        self.to_def(src, keys::BLOCK)
+    }
     pub(super) fn enum_variant_to_def(
         &mut self,
         src: InFile<ast::Variant>,
diff --git a/crates/ide/src/syntax_highlighting/highlight.rs b/crates/ide/src/syntax_highlighting/highlight.rs
index e7c1b4497e2..96c7c475594 100644
--- a/crates/ide/src/syntax_highlighting/highlight.rs
+++ b/crates/ide/src/syntax_highlighting/highlight.rs
@@ -342,9 +342,11 @@ fn highlight_name(
 
 fn calc_binding_hash(name: &hir::Name, shadow_count: u32) -> u64 {
     fn hash<T: std::hash::Hash + std::fmt::Debug>(x: T) -> u64 {
-        use std::{collections::hash_map::DefaultHasher, hash::Hasher};
+        use ide_db::FxHasher;
 
-        let mut hasher = DefaultHasher::new();
+        use std::hash::Hasher;
+
+        let mut hasher = FxHasher::default();
         x.hash(&mut hasher);
         hasher.finish()
     }
diff --git a/crates/ide/src/syntax_highlighting/test_data/highlight_block_mod_items.html b/crates/ide/src/syntax_highlighting/test_data/highlight_block_mod_items.html
new file mode 100644
index 00000000000..977d18c6b73
--- /dev/null
+++ b/crates/ide/src/syntax_highlighting/test_data/highlight_block_mod_items.html
@@ -0,0 +1,64 @@
+
+<style>
+body                { margin: 0; }
+pre                 { color: #DCDCCC; background: #3F3F3F; font-size: 22px; padding: 0.4em; }
+
+.lifetime           { color: #DFAF8F; font-style: italic; }
+.label              { color: #DFAF8F; font-style: italic; }
+.comment            { color: #7F9F7F; }
+.documentation      { color: #629755; }
+.intra_doc_link     { font-style: italic; }
+.injected           { opacity: 0.65 ; }
+.struct, .enum      { color: #7CB8BB; }
+.enum_variant       { color: #BDE0F3; }
+.string_literal     { color: #CC9393; }
+.field              { color: #94BFF3; }
+.function           { color: #93E0E3; }
+.function.unsafe    { color: #BC8383; }
+.trait.unsafe       { color: #BC8383; }
+.operator.unsafe    { color: #BC8383; }
+.mutable.unsafe     { color: #BC8383; text-decoration: underline; }
+.keyword.unsafe     { color: #BC8383; font-weight: bold; }
+.macro.unsafe       { color: #BC8383; }
+.parameter          { color: #94BFF3; }
+.text               { color: #DCDCCC; }
+.type               { color: #7CB8BB; }
+.builtin_type       { color: #8CD0D3; }
+.type_param         { color: #DFAF8F; }
+.attribute          { color: #94BFF3; }
+.numeric_literal    { color: #BFEBBF; }
+.bool_literal       { color: #BFE6EB; }
+.macro              { color: #94BFF3; }
+.derive             { color: #94BFF3; font-style: italic; }
+.module             { color: #AFD8AF; }
+.value_param        { color: #DCDCCC; }
+.variable           { color: #DCDCCC; }
+.format_specifier   { color: #CC696B; }
+.mutable            { text-decoration: underline; }
+.escape_sequence    { color: #94BFF3; }
+.keyword            { color: #F0DFAF; font-weight: bold; }
+.control            { font-style: italic; }
+.reference          { font-style: italic; font-weight: bold; }
+
+.invalid_escape_sequence { color: #FC5555; text-decoration: wavy underline; }
+.unresolved_reference    { color: #FC5555; text-decoration: wavy underline; }
+</style>
+<pre><code><span class="keyword">macro_rules</span><span class="macro_bang">!</span> <span class="macro declaration">foo</span> <span class="brace">{</span>
+    <span class="parenthesis">(</span><span class="punctuation">$</span>foo<span class="colon">:</span>ident<span class="parenthesis">)</span> <span class="operator">=</span><span class="angle">&gt;</span> <span class="brace">{</span>
+        <span class="keyword">mod</span> y <span class="brace">{</span>
+            <span class="keyword">struct</span> <span class="punctuation">$</span>foo<span class="semicolon">;</span>
+        <span class="brace">}</span>
+    <span class="brace">}</span><span class="semicolon">;</span>
+<span class="brace">}</span>
+<span class="keyword">fn</span> <span class="function declaration">main</span><span class="parenthesis">(</span><span class="parenthesis">)</span> <span class="brace">{</span>
+    <span class="macro">foo</span><span class="macro_bang">!</span><span class="parenthesis macro">(</span><span class="struct declaration macro">Foo</span><span class="parenthesis macro">)</span><span class="semicolon">;</span>
+    <span class="keyword">mod</span> <span class="module declaration">module</span> <span class="brace">{</span>
+        <span class="comment">// FIXME: IDE layer has this unresolved</span>
+        <span class="unresolved_reference">foo</span><span class="macro_bang">!</span><span class="parenthesis macro">(</span><span class="none macro">Bar</span><span class="parenthesis macro">)</span><span class="semicolon">;</span>
+        <span class="keyword">fn</span> <span class="function declaration">func</span><span class="parenthesis">(</span><span class="parenthesis">)</span> <span class="brace">{</span>
+            <span class="keyword">mod</span> <span class="module declaration">inner</span> <span class="brace">{</span>
+                <span class="keyword">struct</span> <span class="struct declaration">Innerest</span><span class="angle">&lt;</span><span class="keyword">const</span> <span class="const_param declaration">C</span><span class="colon">:</span> <span class="unresolved_reference">usize</span><span class="angle">&gt;</span> <span class="brace">{</span> <span class="field declaration">field</span><span class="colon">:</span> <span class="bracket">[</span><span class="parenthesis">(</span><span class="parenthesis">)</span><span class="semicolon">;</span> <span class="brace">{</span><span class="const_param">C</span><span class="brace">}</span><span class="bracket">]</span> <span class="brace">}</span>
+            <span class="brace">}</span>
+        <span class="brace">}</span>
+    <span class="brace">}</span>
+<span class="brace">}</span></code></pre>
\ No newline at end of file
diff --git a/crates/ide/src/syntax_highlighting/test_data/highlight_rainbow.html b/crates/ide/src/syntax_highlighting/test_data/highlight_rainbow.html
index ec18c3ea1f9..7ee7b338c19 100644
--- a/crates/ide/src/syntax_highlighting/test_data/highlight_rainbow.html
+++ b/crates/ide/src/syntax_highlighting/test_data/highlight_rainbow.html
@@ -44,14 +44,14 @@ pre                 { color: #DCDCCC; background: #3F3F3F; font-size: 22px; padd
 .unresolved_reference    { color: #FC5555; text-decoration: wavy underline; }
 </style>
 <pre><code><span class="keyword">fn</span> <span class="function declaration">main</span><span class="parenthesis">(</span><span class="parenthesis">)</span> <span class="brace">{</span>
-    <span class="keyword">let</span> <span class="variable declaration reference" data-binding-hash="8121853618659664005" style="color: hsl(273,88%,88%);">hello</span> <span class="operator">=</span> <span class="string_literal">"hello"</span><span class="semicolon">;</span>
-    <span class="keyword">let</span> <span class="variable declaration" data-binding-hash="2705725358298919760" style="color: hsl(76,47%,83%);">x</span> <span class="operator">=</span> <span class="variable reference" data-binding-hash="8121853618659664005" style="color: hsl(273,88%,88%);">hello</span><span class="operator">.</span><span class="unresolved_reference">to_string</span><span class="parenthesis">(</span><span class="parenthesis">)</span><span class="semicolon">;</span>
-    <span class="keyword">let</span> <span class="variable declaration" data-binding-hash="3365759661443752373" style="color: hsl(15,86%,51%);">y</span> <span class="operator">=</span> <span class="variable reference" data-binding-hash="8121853618659664005" style="color: hsl(273,88%,88%);">hello</span><span class="operator">.</span><span class="unresolved_reference">to_string</span><span class="parenthesis">(</span><span class="parenthesis">)</span><span class="semicolon">;</span>
+    <span class="keyword">let</span> <span class="variable declaration reference" data-binding-hash="8384512769119783714" style="color: hsl(59,93%,58%);">hello</span> <span class="operator">=</span> <span class="string_literal">"hello"</span><span class="semicolon">;</span>
+    <span class="keyword">let</span> <span class="variable declaration" data-binding-hash="17360984456076382725" style="color: hsl(95,79%,86%);">x</span> <span class="operator">=</span> <span class="variable reference" data-binding-hash="8384512769119783714" style="color: hsl(59,93%,58%);">hello</span><span class="operator">.</span><span class="unresolved_reference">to_string</span><span class="parenthesis">(</span><span class="parenthesis">)</span><span class="semicolon">;</span>
+    <span class="keyword">let</span> <span class="variable declaration" data-binding-hash="17186414787327620935" style="color: hsl(196,64%,89%);">y</span> <span class="operator">=</span> <span class="variable reference" data-binding-hash="8384512769119783714" style="color: hsl(59,93%,58%);">hello</span><span class="operator">.</span><span class="unresolved_reference">to_string</span><span class="parenthesis">(</span><span class="parenthesis">)</span><span class="semicolon">;</span>
 
-    <span class="keyword">let</span> <span class="variable declaration reference" data-binding-hash="794745962933817518" style="color: hsl(127,71%,87%);">x</span> <span class="operator">=</span> <span class="string_literal">"other color please!"</span><span class="semicolon">;</span>
-    <span class="keyword">let</span> <span class="variable declaration" data-binding-hash="6717528807933952652" style="color: hsl(90,74%,79%);">y</span> <span class="operator">=</span> <span class="variable reference" data-binding-hash="794745962933817518" style="color: hsl(127,71%,87%);">x</span><span class="operator">.</span><span class="unresolved_reference">to_string</span><span class="parenthesis">(</span><span class="parenthesis">)</span><span class="semicolon">;</span>
+    <span class="keyword">let</span> <span class="variable declaration reference" data-binding-hash="4786021388930833562" style="color: hsl(137,61%,87%);">x</span> <span class="operator">=</span> <span class="string_literal">"other color please!"</span><span class="semicolon">;</span>
+    <span class="keyword">let</span> <span class="variable declaration" data-binding-hash="18017815841345165192" style="color: hsl(39,76%,89%);">y</span> <span class="operator">=</span> <span class="variable reference" data-binding-hash="4786021388930833562" style="color: hsl(137,61%,87%);">x</span><span class="operator">.</span><span class="unresolved_reference">to_string</span><span class="parenthesis">(</span><span class="parenthesis">)</span><span class="semicolon">;</span>
 <span class="brace">}</span>
 
 <span class="keyword">fn</span> <span class="function declaration">bar</span><span class="parenthesis">(</span><span class="parenthesis">)</span> <span class="brace">{</span>
-    <span class="keyword">let</span> <span class="keyword">mut</span> <span class="variable declaration mutable reference" data-binding-hash="8121853618659664005" style="color: hsl(273,88%,88%);">hello</span> <span class="operator">=</span> <span class="string_literal">"hello"</span><span class="semicolon">;</span>
+    <span class="keyword">let</span> <span class="keyword">mut</span> <span class="variable declaration mutable reference" data-binding-hash="8384512769119783714" style="color: hsl(59,93%,58%);">hello</span> <span class="operator">=</span> <span class="string_literal">"hello"</span><span class="semicolon">;</span>
 <span class="brace">}</span></code></pre>
\ No newline at end of file
diff --git a/crates/ide/src/syntax_highlighting/tests.rs b/crates/ide/src/syntax_highlighting/tests.rs
index 864c6d1cad7..6fed7d783e8 100644
--- a/crates/ide/src/syntax_highlighting/tests.rs
+++ b/crates/ide/src/syntax_highlighting/tests.rs
@@ -993,10 +993,6 @@ pub struct Struct;
 }
 
 #[test]
-#[cfg_attr(
-    not(all(unix, target_pointer_width = "64")),
-    ignore = "depends on `DefaultHasher` outputs"
-)]
 fn test_rainbow_highlighting() {
     check_highlighting(
         r#"
@@ -1019,6 +1015,35 @@ fn bar() {
 }
 
 #[test]
+fn test_block_mod_items() {
+    check_highlighting(
+        r#"
+macro_rules! foo {
+    ($foo:ident) => {
+        mod y {
+            struct $foo;
+        }
+    };
+}
+fn main() {
+    foo!(Foo);
+    mod module {
+        // FIXME: IDE layer has this unresolved
+        foo!(Bar);
+        fn func() {
+            mod inner {
+                struct Innerest<const C: usize> { field: [(); {C}] }
+            }
+        }
+    }
+}
+"#,
+        expect_file!["./test_data/highlight_block_mod_items.html"],
+        false,
+    );
+}
+
+#[test]
 fn test_ranges() {
     let (analysis, file_id) = fixture::file(
         r#"