Change item collection to be on demand

4 files changed, 38 insertions, 33 deletions

diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index 044d6a2419f..3eaff9c051f 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -3,10 +3,13 @@
 //! For that, we define APIs that will temporarily be public to 3P that exposes rustc internal APIs
 //! until stable MIR is complete.
 
-use crate::stable_mir::CrateItem;
+use crate::stable_mir;
+pub use rustc_span::def_id::{CrateNum, DefId};
 
-pub type DefId = rustc_span::def_id::DefId;
-
-pub fn item_def_id(item: &CrateItem) -> DefId {
+pub fn item_def_id(item: &stable_mir::CrateItem) -> DefId {
     item.0
 }
+
+pub fn crate_num(item: &stable_mir::Crate) -> CrateNum {
+    item.id.into()
+}
diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index 768b01f526f..d956f0ac802 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -9,7 +9,7 @@
 
 use crate::stable_mir::{self};
 use rustc_middle::ty::{tls::with, TyCtxt};
-use rustc_span::def_id::{CrateNum, DefId, CRATE_DEF_INDEX, LOCAL_CRATE};
+use rustc_span::def_id::{CrateNum, LOCAL_CRATE};
 use tracing::debug;
 
 /// Get information about the local crate.
@@ -17,6 +17,11 @@ pub fn local_crate() -> stable_mir::Crate {
     with(|tcx| smir_crate(tcx, LOCAL_CRATE))
 }
 
+/// Retrieve a list of all external crates.
+pub fn external_crates() -> Vec<stable_mir::Crate> {
+    with(|tcx| tcx.crates(()).iter().map(|crate_num| smir_crate(tcx, *crate_num)).collect())
+}
+
 /// Find a crate with the given name.
 pub fn find_crate(name: &str) -> Option<stable_mir::Crate> {
     with(|tcx| {
@@ -27,26 +32,17 @@ pub fn find_crate(name: &str) -> Option<stable_mir::Crate> {
     })
 }
 
+/// Retrieve all items of the local crate that have a MIR associated with them.
+pub fn all_local_items() -> stable_mir::CrateItems {
+    with(|tcx| {
+        tcx.mir_keys(()).iter().map(|item| stable_mir::CrateItem(item.to_def_id())).collect()
+    })
+}
+
 /// Build a stable mir crate from a given crate number.
 fn smir_crate(tcx: TyCtxt<'_>, crate_num: CrateNum) -> stable_mir::Crate {
     let crate_name = tcx.crate_name(crate_num).to_string();
     let is_local = crate_num == LOCAL_CRATE;
-    let mod_id = DefId { index: CRATE_DEF_INDEX, krate: crate_num };
-    let items = if is_local {
-        tcx.hir_module_items(mod_id.expect_local())
-            .items()
-            .map(|item| {
-                let def_id = item.owner_id.def_id.to_def_id();
-                stable_mir::CrateItem(def_id)
-            })
-            .collect()
-    } else {
-        tcx.module_children(mod_id)
-            .iter()
-            .filter_map(|item| item.res.opt_def_id())
-            .map(stable_mir::CrateItem)
-            .collect::<Vec<_>>()
-    };
     debug!(?crate_name, ?crate_num, "smir_crate");
-    stable_mir::Crate { id: crate_num.into(), name: crate_name, is_local, root_items: items }
+    stable_mir::Crate { id: crate_num.into(), name: crate_name, is_local }
 }
diff --git a/compiler/rustc_smir/src/stable_mir/mod.rs b/compiler/rustc_smir/src/stable_mir/mod.rs
index f61ea461dfa..cbf52e691fb 100644
--- a/compiler/rustc_smir/src/stable_mir/mod.rs
+++ b/compiler/rustc_smir/src/stable_mir/mod.rs
@@ -31,8 +31,6 @@ pub struct Crate {
     pub(crate) id: CrateNum,
     pub name: Symbol,
     pub is_local: bool,
-    /// The items defined in the root of this crate.
-    pub root_items: CrateItems,
 }
 
 /// Holds information about an item in the crate.
@@ -50,3 +48,13 @@ pub fn local_crate() -> Crate {
 pub fn find_crate(name: &str) -> Option<Crate> {
     crate::rustc_smir::find_crate(name)
 }
+
+/// Try to find a crate with the given name.
+pub fn external_crates() -> Vec<Crate> {
+    crate::rustc_smir::external_crates()
+}
+
+/// Retrieve all items in the local crate that have a MIR associated with them.
+pub fn all_local_items() -> CrateItems {
+    crate::rustc_smir::all_local_items()
+}
diff --git a/tests/ui-fulldeps/stable-mir/crate-info.rs b/tests/ui-fulldeps/stable-mir/crate-info.rs
index 62cfb809fcb..4458ab0162e 100644
--- a/tests/ui-fulldeps/stable-mir/crate-info.rs
+++ b/tests/ui-fulldeps/stable-mir/crate-info.rs
@@ -29,19 +29,17 @@ fn test_stable_mir(tcx: TyCtxt<'_>) {
     assert_eq!(&local.name, CRATE_NAME);
 
     // Find items in the local crate.
-    assert!(has_root_item(tcx, &local, (DefKind::Fn, "foo_bar")));
-    assert!(has_root_item(tcx, &local, (DefKind::Mod, "foo")));
-    assert!(!has_root_item(tcx, &local, (DefKind::Fn, "foo::bar")));
+    let items = stable_mir::all_local_items();
+    assert!(has_item(tcx, &items, (DefKind::Fn, "foo_bar")));
+    assert!(has_item(tcx, &items, (DefKind::Fn, "foo::bar")));
 
-    // Check that we can find items in the `std` crate.
-    let std_crate = stable_mir::find_crate("std").unwrap();
-    assert!(has_root_item(tcx, &std_crate, (DefKind::Mod, "std::any")));
-    assert!(!has_root_item(tcx, &std_crate, (DefKind::Fn, "std::any::type_name")));
+    // Find the `std` crate.
+    assert!(stable_mir::find_crate("std").is_some());
 }
 
 // Use internal API to find a function in a crate.
-fn has_root_item(tcx: TyCtxt, krate: &stable_mir::Crate, item: (DefKind, &str)) -> bool {
-    krate.root_items.iter().any(|crate_item| {
+fn has_item(tcx: TyCtxt, items: &stable_mir::CrateItems, item: (DefKind, &str)) -> bool {
+    items.iter().any(|crate_item| {
         let def_id = rustc_internal::item_def_id(crate_item);
         tcx.def_kind(def_id) == item.0 && tcx.def_path_str(def_id) == item.1
     })