Gather module items after lowering.

12 files changed, 119 insertions, 63 deletions

diff --git a/compiler/rustc_ast_lowering/src/item.rs b/compiler/rustc_ast_lowering/src/item.rs
index b7497c713f3..252bc676dbe 100644
--- a/compiler/rustc_ast_lowering/src/item.rs
+++ b/compiler/rustc_ast_lowering/src/item.rs
@@ -50,12 +50,6 @@ impl<'a> Visitor<'a> for ItemLowerer<'a, '_, '_> {
         self.lctx.with_parent_item_lifetime_defs(hir_id, |this| {
             let this = &mut ItemLowerer { lctx: this };
             match item.kind {
-                ItemKind::Mod(..) => {
-                    let def_id = this.lctx.lower_node_id(item.id).expect_owner();
-                    let old_current_module = mem::replace(&mut this.lctx.current_module, def_id);
-                    visit::walk_item(this, item);
-                    this.lctx.current_module = old_current_module;
-                }
                 ItemKind::Impl(box ImplKind { ref of_trait, .. }) => {
                     this.with_trait_impl_ref(of_trait, |this| visit::walk_item(this, item));
                 }
diff --git a/compiler/rustc_ast_lowering/src/lib.rs b/compiler/rustc_ast_lowering/src/lib.rs
index 8d731d7a578..5cda9c4aae3 100644
--- a/compiler/rustc_ast_lowering/src/lib.rs
+++ b/compiler/rustc_ast_lowering/src/lib.rs
@@ -102,8 +102,6 @@ struct LoweringContext<'a, 'hir: 'a> {
     owners: IndexVec<LocalDefId, Option<hir::OwnerNode<'hir>>>,
     bodies: BTreeMap<hir::BodyId, hir::Body<'hir>>,
 
-    modules: BTreeMap<LocalDefId, hir::ModuleItems>,
-
     generator_kind: Option<hir::GeneratorKind>,
 
     attrs: BTreeMap<hir::HirId, &'hir [Attribute]>,
@@ -152,8 +150,6 @@ struct LoweringContext<'a, 'hir: 'a> {
     /// vector.
     in_scope_lifetimes: Vec<ParamName>,
 
-    current_module: LocalDefId,
-
     current_hir_id_owner: (LocalDefId, u32),
     item_local_id_counters: NodeMap<u32>,
     node_id_to_hir_id: IndexVec<NodeId, Option<hir::HirId>>,
@@ -327,7 +323,6 @@ pub fn lower_crate<'a, 'hir>(
         arena,
         owners: IndexVec::default(),
         bodies: BTreeMap::new(),
-        modules: BTreeMap::new(),
         attrs: BTreeMap::default(),
         catch_scope: None,
         loop_scope: None,
@@ -335,7 +330,6 @@ pub fn lower_crate<'a, 'hir>(
         is_in_trait_impl: false,
         is_in_dyn_type: false,
         anonymous_lifetime_mode: AnonymousLifetimeMode::PassThrough,
-        current_module: CRATE_DEF_ID,
         current_hir_id_owner: (CRATE_DEF_ID, 0),
         item_local_id_counters: Default::default(),
         node_id_to_hir_id: IndexVec::new(),
@@ -508,13 +502,8 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             }
         }
 
-        let krate = hir::Crate {
-            owners: self.owners,
-            bodies: self.bodies,
-            modules: self.modules,
-            trait_map,
-            attrs: self.attrs,
-        };
+        let krate =
+            hir::Crate { owners: self.owners, bodies: self.bodies, trait_map, attrs: self.attrs };
         self.arena.alloc(krate)
     }
 
@@ -523,7 +512,6 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         let item = self.arena.alloc(item);
         self.owners.ensure_contains_elem(id.def_id, || None);
         self.owners[id.def_id] = Some(hir::OwnerNode::Item(item));
-        self.modules.entry(self.current_module).or_default().items.insert(id);
         id
     }
 
@@ -532,7 +520,6 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         let item = self.arena.alloc(item);
         self.owners.ensure_contains_elem(id.def_id, || None);
         self.owners[id.def_id] = Some(hir::OwnerNode::ForeignItem(item));
-        self.modules.entry(self.current_module).or_default().foreign_items.insert(id);
         id
     }
 
@@ -541,7 +528,6 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         let item = self.arena.alloc(item);
         self.owners.ensure_contains_elem(id.def_id, || None);
         self.owners[id.def_id] = Some(hir::OwnerNode::ImplItem(item));
-        self.modules.entry(self.current_module).or_default().impl_items.insert(id);
         id
     }
 
@@ -550,7 +536,6 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         let item = self.arena.alloc(item);
         self.owners.ensure_contains_elem(id.def_id, || None);
         self.owners[id.def_id] = Some(hir::OwnerNode::TraitItem(item));
-        self.modules.entry(self.current_module).or_default().trait_items.insert(id);
         id
     }
 
diff --git a/compiler/rustc_hir/src/hir.rs b/compiler/rustc_hir/src/hir.rs
index 84bc37170c6..7ed74375c51 100644
--- a/compiler/rustc_hir/src/hir.rs
+++ b/compiler/rustc_hir/src/hir.rs
@@ -21,7 +21,7 @@ use rustc_target::asm::InlineAsmRegOrRegClass;
 use rustc_target::spec::abi::Abi;
 
 use smallvec::SmallVec;
-use std::collections::{BTreeMap, BTreeSet};
+use std::collections::BTreeMap;
 use std::fmt;
 
 #[derive(Copy, Clone, Encodable, HashStable_Generic)]
@@ -653,16 +653,6 @@ pub struct WhereEqPredicate<'hir> {
     pub rhs_ty: &'hir Ty<'hir>,
 }
 
-#[derive(Default, Encodable, Debug, HashStable_Generic)]
-pub struct ModuleItems {
-    // Use BTreeSets here so items are in the same order as in the
-    // list of all items in Crate
-    pub items: BTreeSet<ItemId>,
-    pub trait_items: BTreeSet<TraitItemId>,
-    pub impl_items: BTreeSet<ImplItemId>,
-    pub foreign_items: BTreeSet<ForeignItemId>,
-}
-
 /// The top-level data structure that stores the entire contents of
 /// the crate currently being compiled.
 ///
@@ -674,10 +664,6 @@ pub struct Crate<'hir> {
     pub owners: IndexVec<LocalDefId, Option<OwnerNode<'hir>>>,
     pub bodies: BTreeMap<BodyId, Body<'hir>>,
 
-    /// A list of modules written out in the order in which they
-    /// appear in the crate. This includes the main crate module.
-    pub modules: BTreeMap<LocalDefId, ModuleItems>,
-
     /// Map indicating what traits are in scope for places where this
     /// is relevant; generated by resolve.
     pub trait_map: FxHashMap<LocalDefId, FxHashMap<ItemLocalId, Box<[TraitCandidate]>>>,
diff --git a/compiler/rustc_interface/src/passes.rs b/compiler/rustc_interface/src/passes.rs
index 1f3d6f70ff8..cb7529b527e 100644
--- a/compiler/rustc_interface/src/passes.rs
+++ b/compiler/rustc_interface/src/passes.rs
@@ -8,7 +8,7 @@ use rustc_borrowck as mir_borrowck;
 use rustc_codegen_ssa::back::link::emit_metadata;
 use rustc_codegen_ssa::traits::CodegenBackend;
 use rustc_data_structures::parallel;
-use rustc_data_structures::sync::{par_iter, Lrc, OnceCell, ParallelIterator, WorkerLocal};
+use rustc_data_structures::sync::{Lrc, OnceCell, WorkerLocal};
 use rustc_data_structures::temp_dir::MaybeTempDir;
 use rustc_errors::{ErrorReported, PResult};
 use rustc_expand::base::ExtCtxt;
@@ -861,7 +861,7 @@ fn analysis(tcx: TyCtxt<'_>, (): ()) -> Result<()> {
                 CStore::from_tcx(tcx).report_unused_deps(tcx);
             },
             {
-                par_iter(&tcx.hir().krate().modules).for_each(|(&module, _)| {
+                tcx.hir().par_for_each_module(|module| {
                     tcx.ensure().check_mod_loops(module);
                     tcx.ensure().check_mod_attrs(module);
                     tcx.ensure().check_mod_naked_functions(module);
@@ -893,7 +893,7 @@ fn analysis(tcx: TyCtxt<'_>, (): ()) -> Result<()> {
             },
             {
                 sess.time("liveness_and_intrinsic_checking", || {
-                    par_iter(&tcx.hir().krate().modules).for_each(|(&module, _)| {
+                    tcx.hir().par_for_each_module(|module| {
                         // this must run before MIR dump, because
                         // "not all control paths return a value" is reported here.
                         //
@@ -963,7 +963,7 @@ fn analysis(tcx: TyCtxt<'_>, (): ()) -> Result<()> {
             },
             {
                 sess.time("privacy_checking_modules", || {
-                    par_iter(&tcx.hir().krate().modules).for_each(|(&module, _)| {
+                    tcx.hir().par_for_each_module(|module| {
                         tcx.ensure().check_mod_privacy(module);
                     });
                 });
diff --git a/compiler/rustc_lint/src/late.rs b/compiler/rustc_lint/src/late.rs
index 2070fd69d3f..00c3a6fa25e 100644
--- a/compiler/rustc_lint/src/late.rs
+++ b/compiler/rustc_lint/src/late.rs
@@ -16,7 +16,7 @@
 
 use crate::{passes::LateLintPassObject, LateContext, LateLintPass, LintStore};
 use rustc_ast as ast;
-use rustc_data_structures::sync::{join, par_iter, ParallelIterator};
+use rustc_data_structures::sync::join;
 use rustc_hir as hir;
 use rustc_hir::def_id::LocalDefId;
 use rustc_hir::intravisit as hir_visit;
@@ -501,9 +501,7 @@ pub fn check_crate<'tcx, T: LateLintPass<'tcx>>(
         || {
             tcx.sess.time("module_lints", || {
                 // Run per-module lints
-                par_iter(&tcx.hir().krate().modules).for_each(|(&module, _)| {
-                    tcx.ensure().lint_mod(module);
-                });
+                tcx.hir().par_for_each_module(|module| tcx.ensure().lint_mod(module));
             });
         },
     );
diff --git a/compiler/rustc_middle/src/hir/map/mod.rs b/compiler/rustc_middle/src/hir/map/mod.rs
index 10714a4b706..61ba933c32e 100644
--- a/compiler/rustc_middle/src/hir/map/mod.rs
+++ b/compiler/rustc_middle/src/hir/map/mod.rs
@@ -1,11 +1,12 @@
 use self::collector::NodeCollector;
 
-use crate::hir::{AttributeMap, IndexedHir, Owner};
+use crate::hir::{AttributeMap, IndexedHir, ModuleItems, Owner};
 use crate::ty::TyCtxt;
 use rustc_ast as ast;
 use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::svh::Svh;
+use rustc_data_structures::sync::{self, par_iter};
 use rustc_hir::def::{DefKind, Res};
 use rustc_hir::def_id::{CrateNum, DefId, LocalDefId, CRATE_DEF_ID, CRATE_DEF_INDEX, LOCAL_CRATE};
 use rustc_hir::definitions::{DefKey, DefPath, DefPathHash};
@@ -19,6 +20,7 @@ use rustc_span::source_map::Spanned;
 use rustc_span::symbol::{kw, sym, Ident, Symbol};
 use rustc_span::Span;
 use rustc_target::spec::abi::Abi;
+use std::collections::BTreeSet;
 
 pub mod blocks;
 mod collector;
@@ -558,6 +560,31 @@ impl<'hir> Map<'hir> {
         }
     }
 
+    pub fn for_each_module(&self, f: impl Fn(LocalDefId)) {
+        let mut queue = BTreeSet::default();
+        queue.insert(CRATE_DEF_ID);
+
+        while let Some(id) = queue.pop_first() {
+            f(id);
+            let items = self.tcx.hir_module_items(id);
+            queue.extend(items.submodules.iter().copied())
+        }
+    }
+
+    pub fn par_for_each_module(&self, f: impl Fn(LocalDefId) + sync::Sync) {
+        use rustc_data_structures::sync::ParallelIterator;
+        par_iter_submodules(self.tcx, CRATE_DEF_ID, &f);
+
+        fn par_iter_submodules<F>(tcx: TyCtxt<'_>, module: LocalDefId, f: &F)
+        where
+            F: Fn(LocalDefId) + sync::Sync,
+        {
+            (*f)(module);
+            let items = tcx.hir_module_items(module);
+            par_iter(&items.submodules).for_each(|&sm| par_iter_submodules(tcx, sm, f));
+        }
+    }
+
     /// Returns an iterator for the nodes in the ancestor tree of the `current_id`
     /// until the crate root is reached. Prefer this over your own loop using `get_parent_node`.
     pub fn parent_iter(&self, current_id: HirId) -> ParentHirIterator<'_, 'hir> {
@@ -1118,3 +1145,63 @@ fn hir_id_to_string(map: &Map<'_>, id: HirId) -> String {
         None => format!("unknown node{}", id_str),
     }
 }
+
+pub(super) fn hir_module_items(tcx: TyCtxt<'_>, module_id: LocalDefId) -> ModuleItems {
+    let mut collector = ModuleCollector {
+        tcx,
+        submodules: BTreeSet::default(),
+        items: BTreeSet::default(),
+        trait_items: BTreeSet::default(),
+        impl_items: BTreeSet::default(),
+        foreign_items: BTreeSet::default(),
+    };
+
+    let (hir_mod, span, hir_id) = tcx.hir().get_module(module_id);
+    collector.visit_mod(hir_mod, span, hir_id);
+
+    let ModuleCollector { submodules, items, trait_items, impl_items, foreign_items, .. } =
+        collector;
+    return ModuleItems { submodules, items, trait_items, impl_items, foreign_items };
+
+    struct ModuleCollector<'tcx> {
+        tcx: TyCtxt<'tcx>,
+        submodules: BTreeSet<LocalDefId>,
+        items: BTreeSet<ItemId>,
+        trait_items: BTreeSet<TraitItemId>,
+        impl_items: BTreeSet<ImplItemId>,
+        foreign_items: BTreeSet<ForeignItemId>,
+    }
+
+    impl<'hir> Visitor<'hir> for ModuleCollector<'hir> {
+        type Map = Map<'hir>;
+
+        fn nested_visit_map(&mut self) -> intravisit::NestedVisitorMap<Self::Map> {
+            intravisit::NestedVisitorMap::All(self.tcx.hir())
+        }
+
+        fn visit_item(&mut self, item: &'hir Item<'hir>) {
+            self.items.insert(item.item_id());
+            if let ItemKind::Mod(..) = item.kind {
+                // If this declares another module, do not recurse inside it.
+                self.submodules.insert(item.def_id);
+            } else {
+                intravisit::walk_item(self, item)
+            }
+        }
+
+        fn visit_trait_item(&mut self, item: &'hir TraitItem<'hir>) {
+            self.trait_items.insert(item.trait_item_id());
+            intravisit::walk_trait_item(self, item)
+        }
+
+        fn visit_impl_item(&mut self, item: &'hir ImplItem<'hir>) {
+            self.impl_items.insert(item.impl_item_id());
+            intravisit::walk_impl_item(self, item)
+        }
+
+        fn visit_foreign_item(&mut self, item: &'hir ForeignItem<'hir>) {
+            self.foreign_items.insert(item.foreign_item_id());
+            intravisit::walk_foreign_item(self, item)
+        }
+    }
+}
diff --git a/compiler/rustc_middle/src/hir/mod.rs b/compiler/rustc_middle/src/hir/mod.rs
index c8ea7454f0b..7d660b18ae1 100644
--- a/compiler/rustc_middle/src/hir/mod.rs
+++ b/compiler/rustc_middle/src/hir/mod.rs
@@ -17,7 +17,7 @@ use rustc_hir::def_id::LocalDefId;
 use rustc_hir::*;
 use rustc_index::vec::{Idx, IndexVec};
 use rustc_span::DUMMY_SP;
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 
 /// Result of HIR indexing.
 #[derive(Debug)]
@@ -121,6 +121,17 @@ impl<'tcx> AttributeMap<'tcx> {
     }
 }
 
+#[derive(Default, Encodable, Debug, HashStable)]
+pub struct ModuleItems {
+    // Use BTreeSets here so items are in the same order as in the
+    // list of all items in Crate
+    submodules: BTreeSet<LocalDefId>,
+    items: BTreeSet<ItemId>,
+    trait_items: BTreeSet<TraitItemId>,
+    impl_items: BTreeSet<ImplItemId>,
+    foreign_items: BTreeSet<ForeignItemId>,
+}
+
 impl<'tcx> TyCtxt<'tcx> {
     #[inline(always)]
     pub fn hir(self) -> map::Map<'tcx> {
@@ -140,7 +151,7 @@ pub fn provide(providers: &mut Providers) {
     providers.hir_crate = |tcx, ()| tcx.untracked_crate;
     providers.index_hir = map::index_hir;
     providers.crate_hash = map::crate_hash;
-    providers.hir_module_items = |tcx, id| &tcx.untracked_crate.modules[&id];
+    providers.hir_module_items = map::hir_module_items;
     providers.hir_owner = |tcx, id| {
         let owner = tcx.index_hir(()).map[id].as_ref()?;
         let node = owner.nodes[ItemLocalId::new(0)].as_ref().unwrap().node;
diff --git a/compiler/rustc_middle/src/lib.rs b/compiler/rustc_middle/src/lib.rs
index 94267ec64c7..02f0294c8ad 100644
--- a/compiler/rustc_middle/src/lib.rs
+++ b/compiler/rustc_middle/src/lib.rs
@@ -33,6 +33,7 @@
 #![feature(discriminant_kind)]
 #![feature(exhaustive_patterns)]
 #![feature(if_let_guard)]
+#![feature(map_first_last)]
 #![feature(never_type)]
 #![feature(extern_types)]
 #![feature(new_uninit)]
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index d6f3b6f3248..0ff444faa7a 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -52,8 +52,8 @@ rustc_queries! {
     ///
     /// This can be conveniently accessed by `tcx.hir().visit_item_likes_in_module`.
     /// Avoid calling this query directly.
-    query hir_module_items(key: LocalDefId) -> &'tcx hir::ModuleItems {
-        eval_always
+    query hir_module_items(key: LocalDefId) -> rustc_middle::hir::ModuleItems {
+        storage(ArenaCacheSelector<'tcx>)
         desc { |tcx| "HIR module items in `{}`", tcx.def_path_str(key.to_def_id()) }
     }
 
diff --git a/compiler/rustc_passes/src/hir_id_validator.rs b/compiler/rustc_passes/src/hir_id_validator.rs
index 18f61c6e1c1..eff1096c855 100644
--- a/compiler/rustc_passes/src/hir_id_validator.rs
+++ b/compiler/rustc_passes/src/hir_id_validator.rs
@@ -1,5 +1,5 @@
 use rustc_data_structures::fx::FxHashSet;
-use rustc_data_structures::sync::{par_iter, Lock, ParallelIterator};
+use rustc_data_structures::sync::Lock;
 use rustc_hir as hir;
 use rustc_hir::def_id::{LocalDefId, CRATE_DEF_INDEX};
 use rustc_hir::intravisit;
@@ -18,9 +18,9 @@ pub fn check_crate(tcx: TyCtxt<'_>) {
     let errors = Lock::new(Vec::new());
     let hir_map = tcx.hir();
 
-    par_iter(&hir_map.krate().modules).for_each(|(&module_id, _)| {
+    hir_map.par_for_each_module(|module_id| {
         hir_map
-            .visit_item_likes_in_module(module_id, &mut OuterVisitor { hir_map, errors: &errors });
+            .visit_item_likes_in_module(module_id, &mut OuterVisitor { hir_map, errors: &errors })
     });
 
     let errors = errors.into_inner();
diff --git a/compiler/rustc_typeck/src/impl_wf_check.rs b/compiler/rustc_typeck/src/impl_wf_check.rs
index 194c4efdbb0..9b23bf241cc 100644
--- a/compiler/rustc_typeck/src/impl_wf_check.rs
+++ b/compiler/rustc_typeck/src/impl_wf_check.rs
@@ -58,9 +58,7 @@ pub fn impl_wf_check(tcx: TyCtxt<'_>) {
     // We will tag this as part of the WF check -- logically, it is,
     // but it's one that we must perform earlier than the rest of
     // WfCheck.
-    for &module in tcx.hir().krate().modules.keys() {
-        tcx.ensure().check_mod_impl_wf(module);
-    }
+    tcx.hir().for_each_module(|module| tcx.ensure().check_mod_impl_wf(module))
 }
 
 fn check_mod_impl_wf(tcx: TyCtxt<'_>, module_def_id: LocalDefId) {
diff --git a/compiler/rustc_typeck/src/lib.rs b/compiler/rustc_typeck/src/lib.rs
index c703ca96ada..f8714cdc70c 100644
--- a/compiler/rustc_typeck/src/lib.rs
+++ b/compiler/rustc_typeck/src/lib.rs
@@ -473,9 +473,7 @@ pub fn check_crate(tcx: TyCtxt<'_>) -> Result<(), ErrorReported> {
     // FIXME(matthewjasper) We shouldn't need to use `track_errors`.
     tcx.sess.track_errors(|| {
         tcx.sess.time("type_collecting", || {
-            for &module in tcx.hir().krate().modules.keys() {
-                tcx.ensure().collect_mod_item_types(module);
-            }
+            tcx.hir().for_each_module(|module| tcx.ensure().collect_mod_item_types(module))
         });
     })?;
 
@@ -505,9 +503,7 @@ pub fn check_crate(tcx: TyCtxt<'_>) -> Result<(), ErrorReported> {
 
     // NOTE: This is copy/pasted in librustdoc/core.rs and should be kept in sync.
     tcx.sess.time("item_types_checking", || {
-        for &module in tcx.hir().krate().modules.keys() {
-            tcx.ensure().check_mod_item_types(module);
-        }
+        tcx.hir().for_each_module(|module| tcx.ensure().check_mod_item_types(module))
     });
 
     tcx.sess.time("item_bodies_checking", || tcx.typeck_item_bodies(()));