Gather module items after lowering.

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)
+        }
+    }
+}

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;