resolve: Split `module_map` into two maps for local and extern modules

3 files changed, 69 insertions, 41 deletions

diff --git a/compiler/rustc_resolve/src/build_reduced_graph.rs b/compiler/rustc_resolve/src/build_reduced_graph.rs
index 5a1be618cee..676acd78ea7 100644
--- a/compiler/rustc_resolve/src/build_reduced_graph.rs
+++ b/compiler/rustc_resolve/src/build_reduced_graph.rs
@@ -102,33 +102,36 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
     /// or trait), then this function returns that module's resolver representation, otherwise it
     /// returns `None`.
     pub(crate) fn get_module(&mut self, def_id: DefId) -> Option<Module<'ra>> {
-        if let module @ Some(..) = self.module_map.get(&def_id) {
-            return module.copied();
-        }
+        match def_id.as_local() {
+            Some(local_def_id) => self.local_module_map.get(&local_def_id).copied(),
+            None => {
+                if let module @ Some(..) = self.extern_module_map.borrow().get(&def_id) {
+                    return module.copied();
+                }
 
-        if !def_id.is_local() {
-            // Query `def_kind` is not used because query system overhead is too expensive here.
-            let def_kind = self.cstore().def_kind_untracked(def_id);
-            if def_kind.is_module_like() {
-                let parent = self
-                    .tcx
-                    .opt_parent(def_id)
-                    .map(|parent_id| self.get_nearest_non_block_module(parent_id));
-                // Query `expn_that_defined` is not used because
-                // hashing spans in its result is expensive.
-                let expn_id = self.cstore().expn_that_defined_untracked(def_id, self.tcx.sess);
-                return Some(self.new_module(
-                    parent,
-                    ModuleKind::Def(def_kind, def_id, Some(self.tcx.item_name(def_id))),
-                    expn_id,
-                    self.def_span(def_id),
-                    // FIXME: Account for `#[no_implicit_prelude]` attributes.
-                    parent.is_some_and(|module| module.no_implicit_prelude),
-                ));
+                // Query `def_kind` is not used because query system overhead is too expensive here.
+                let def_kind = self.cstore().def_kind_untracked(def_id);
+                if def_kind.is_module_like() {
+                    let parent = self
+                        .tcx
+                        .opt_parent(def_id)
+                        .map(|parent_id| self.get_nearest_non_block_module(parent_id));
+                    // Query `expn_that_defined` is not used because
+                    // hashing spans in its result is expensive.
+                    let expn_id = self.cstore().expn_that_defined_untracked(def_id, self.tcx.sess);
+                    return Some(self.new_extern_module(
+                        parent,
+                        ModuleKind::Def(def_kind, def_id, Some(self.tcx.item_name(def_id))),
+                        expn_id,
+                        self.def_span(def_id),
+                        // FIXME: Account for `#[no_implicit_prelude]` attributes.
+                        parent.is_some_and(|module| module.no_implicit_prelude),
+                    ));
+                }
+
+                None
             }
         }
-
-        None
     }
 
     pub(crate) fn expn_def_scope(&mut self, expn_id: ExpnId) -> Module<'ra> {
@@ -758,7 +761,7 @@ impl<'a, 'ra, 'tcx> BuildReducedGraphVisitor<'a, 'ra, 'tcx> {
                 if let ast::ModKind::Loaded(_, _, _, Err(_)) = mod_kind {
                     self.r.mods_with_parse_errors.insert(def_id);
                 }
-                self.parent_scope.module = self.r.new_module(
+                self.parent_scope.module = self.r.new_local_module(
                     Some(parent),
                     ModuleKind::Def(def_kind, def_id, Some(ident.name)),
                     expansion.to_expn_id(),
@@ -790,7 +793,7 @@ impl<'a, 'ra, 'tcx> BuildReducedGraphVisitor<'a, 'ra, 'tcx> {
             ItemKind::Enum(ident, _, _) | ItemKind::Trait(box ast::Trait { ident, .. }) => {
                 self.r.define(parent, ident, TypeNS, res, vis, sp, expansion);
 
-                self.parent_scope.module = self.r.new_module(
+                self.parent_scope.module = self.r.new_local_module(
                     Some(parent),
                     ModuleKind::Def(def_kind, def_id, Some(ident.name)),
                     expansion.to_expn_id(),
@@ -986,7 +989,7 @@ impl<'a, 'ra, 'tcx> BuildReducedGraphVisitor<'a, 'ra, 'tcx> {
         let parent = self.parent_scope.module;
         let expansion = self.parent_scope.expansion;
         if self.block_needs_anonymous_module(block) {
-            let module = self.r.new_module(
+            let module = self.r.new_local_module(
                 Some(parent),
                 ModuleKind::Block,
                 expansion.to_expn_id(),
diff --git a/compiler/rustc_resolve/src/diagnostics.rs b/compiler/rustc_resolve/src/diagnostics.rs
index d263ab25045..9e9ab7356b8 100644
--- a/compiler/rustc_resolve/src/diagnostics.rs
+++ b/compiler/rustc_resolve/src/diagnostics.rs
@@ -2159,7 +2159,16 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
             .keys()
             .map(|ident| ident.name)
             .chain(
-                self.module_map
+                self.local_module_map
+                    .iter()
+                    .filter(|(_, module)| {
+                        current_module.is_ancestor_of(**module) && current_module != **module
+                    })
+                    .flat_map(|(_, module)| module.kind.name()),
+            )
+            .chain(
+                self.extern_module_map
+                    .borrow()
                     .iter()
                     .filter(|(_, module)| {
                         current_module.is_ancestor_of(**module) && current_module != **module
diff --git a/compiler/rustc_resolve/src/lib.rs b/compiler/rustc_resolve/src/lib.rs
index 73a6d4c0c3e..98786c1e6f9 100644
--- a/compiler/rustc_resolve/src/lib.rs
+++ b/compiler/rustc_resolve/src/lib.rs
@@ -1081,7 +1081,10 @@ pub struct Resolver<'ra, 'tcx> {
     /// some AST passes can generate identifiers that only resolve to local or
     /// lang items.
     empty_module: Module<'ra>,
-    module_map: FxIndexMap<DefId, Module<'ra>>,
+    /// Eagerly populated map of all local non-block modules.
+    local_module_map: FxIndexMap<LocalDefId, Module<'ra>>,
+    /// Lazily populated cache of modules loaded from external crates.
+    extern_module_map: RefCell<FxIndexMap<DefId, Module<'ra>>>,
     binding_parent_modules: FxHashMap<NameBinding<'ra>, Module<'ra>>,
 
     underscore_disambiguator: u32,
@@ -1112,8 +1115,9 @@ pub struct Resolver<'ra, 'tcx> {
     builtin_macros: FxHashMap<Symbol, SyntaxExtensionKind>,
     registered_tools: &'tcx RegisteredTools,
     macro_use_prelude: FxIndexMap<Symbol, NameBinding<'ra>>,
+    /// Eagerly populated map of all local macro definitions.
     local_macro_map: FxHashMap<LocalDefId, &'ra MacroData>,
-    /// Lazily populated cache of macros loaded from external crates.
+    /// Lazily populated cache of macro definitions loaded from external crates.
     extern_macro_map: RefCell<FxHashMap<DefId, &'ra MacroData>>,
     dummy_ext_bang: Arc<SyntaxExtension>,
     dummy_ext_derive: Arc<SyntaxExtension>,
@@ -1265,7 +1269,6 @@ impl<'ra> ResolverArenas<'ra> {
         expn_id: ExpnId,
         span: Span,
         no_implicit_prelude: bool,
-        module_map: &mut FxIndexMap<DefId, Module<'ra>>,
     ) -> Module<'ra> {
         let (def_id, self_binding) = match kind {
             ModuleKind::Def(def_kind, def_id, _) => (
@@ -1285,9 +1288,6 @@ impl<'ra> ResolverArenas<'ra> {
         if def_id.is_none_or(|def_id| def_id.is_local()) {
             self.local_modules.borrow_mut().push(module);
         }
-        if let Some(def_id) = def_id {
-            module_map.insert(def_id, module);
-        }
         module
     }
     fn local_modules(&'ra self) -> std::cell::Ref<'ra, Vec<Module<'ra>>> {
@@ -1428,22 +1428,21 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
         arenas: &'ra ResolverArenas<'ra>,
     ) -> Resolver<'ra, 'tcx> {
         let root_def_id = CRATE_DEF_ID.to_def_id();
-        let mut module_map = FxIndexMap::default();
+        let mut local_module_map = FxIndexMap::default();
         let graph_root = arenas.new_module(
             None,
             ModuleKind::Def(DefKind::Mod, root_def_id, None),
             ExpnId::root(),
             crate_span,
             attr::contains_name(attrs, sym::no_implicit_prelude),
-            &mut module_map,
         );
+        local_module_map.insert(CRATE_DEF_ID, graph_root);
         let empty_module = arenas.new_module(
             None,
             ModuleKind::Def(DefKind::Mod, root_def_id, None),
             ExpnId::root(),
             DUMMY_SP,
             true,
-            &mut Default::default(),
         );
 
         let mut node_id_to_def_id = NodeMap::default();
@@ -1504,7 +1503,8 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
             trait_map: NodeMap::default(),
             underscore_disambiguator: 0,
             empty_module,
-            module_map,
+            local_module_map,
+            extern_module_map: Default::default(),
             block_map: Default::default(),
             binding_parent_modules: FxHashMap::default(),
             ast_transform_scopes: FxHashMap::default(),
@@ -1608,7 +1608,22 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
         resolver
     }
 
-    fn new_module(
+    fn new_local_module(
+        &mut self,
+        parent: Option<Module<'ra>>,
+        kind: ModuleKind,
+        expn_id: ExpnId,
+        span: Span,
+        no_implicit_prelude: bool,
+    ) -> Module<'ra> {
+        let module = self.arenas.new_module(parent, kind, expn_id, span, no_implicit_prelude);
+        if let Some(def_id) = module.opt_def_id() {
+            self.local_module_map.insert(def_id.expect_local(), module);
+        }
+        module
+    }
+
+    fn new_extern_module(
         &mut self,
         parent: Option<Module<'ra>>,
         kind: ModuleKind,
@@ -1616,8 +1631,9 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
         span: Span,
         no_implicit_prelude: bool,
     ) -> Module<'ra> {
-        let module_map = &mut self.module_map;
-        self.arenas.new_module(parent, kind, expn_id, span, no_implicit_prelude, module_map)
+        let module = self.arenas.new_module(parent, kind, expn_id, span, no_implicit_prelude);
+        self.extern_module_map.borrow_mut().insert(module.def_id(), module);
+        module
     }
 
     fn new_local_macro(&mut self, def_id: LocalDefId, macro_data: MacroData) -> &'ra MacroData {