From 89825439017a7e0a83084e6e089176e9458abccb Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Thu, 7 Nov 2024 11:48:48 -0800
Subject: [PATCH] Improve query for local crate and simplify smir code

---
 compiler/rustc_middle/src/hir/mod.rs          |  2 +-
 compiler/rustc_middle/src/query/mod.rs        |  2 --
 compiler/rustc_smir/src/rustc_smir/context.rs | 24 ++-----------------
 compiler/rustc_smir/src/rustc_smir/mod.rs     |  6 +----
 4 files changed, 4 insertions(+), 30 deletions(-)

diff --git a/compiler/rustc_middle/src/hir/mod.rs b/compiler/rustc_middle/src/hir/mod.rs
index 89e0c4ab85d4..60b82f4d1d81 100644
--- a/compiler/rustc_middle/src/hir/mod.rs
+++ b/compiler/rustc_middle/src/hir/mod.rs
@@ -233,5 +233,5 @@ pub fn provide(providers: &mut Providers) {
     providers.in_scope_traits_map = |tcx, id| {
         tcx.hir_crate(()).owners[id.def_id].as_owner().map(|owner_info| &owner_info.trait_map)
     };
-    providers.num_def_ids = |tcx, _| tcx.hir_crate_items(()).definitions().count();
+    providers.num_def_ids = |tcx, _| tcx.definitions_untracked().num_definitions();
 }
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index f04f6bb7ec1b..a0020c71c50a 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -1815,8 +1815,6 @@ rustc_queries! {
     /// Gets the number of definitions in this crate.
     ///
     /// This allows external tools to iterate over all definitions in a crate.
-    /// For completeness, this also accepts `LOCAL_CRATE` as an argument, but consider using
-    /// `hir_crate_items(())` instead.
     query num_def_ids(_: CrateNum) -> usize {
         desc { "fetching the number of definitions in a crate" }
         separate_provide_extern
diff --git a/compiler/rustc_smir/src/rustc_smir/context.rs b/compiler/rustc_smir/src/rustc_smir/context.rs
index 7c6f80312404..ddcd64b68b4f 100644
--- a/compiler/rustc_smir/src/rustc_smir/context.rs
+++ b/compiler/rustc_smir/src/rustc_smir/context.rs
@@ -84,34 +84,14 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         let mut tables = self.0.borrow_mut();
         let tcx = tables.tcx;
         let krate = crate_num.internal(&mut *tables, tcx);
-        if krate == LOCAL_CRATE {
-            tcx.hir_crate_items(())
-                .definitions()
-                .filter_map(|local_id| tables.to_fn_def(tcx, local_id.to_def_id()))
-                .collect()
-        } else {
-            tables
-                .iter_external_def_id(tcx, krate)
-                .filter_map(|def_id| tables.to_fn_def(tcx, def_id))
-                .collect()
-        }
+        tables.iter_def_ids(tcx, krate).filter_map(|def_id| tables.to_fn_def(tcx, def_id)).collect()
     }
 
     fn crate_statics(&self, crate_num: CrateNum) -> Vec<StaticDef> {
         let mut tables = self.0.borrow_mut();
         let tcx = tables.tcx;
         let krate = crate_num.internal(&mut *tables, tcx);
-        if krate == LOCAL_CRATE {
-            tcx.hir_crate_items(())
-                .definitions()
-                .filter_map(|local_id| tables.to_static(tcx, local_id.to_def_id()))
-                .collect()
-        } else {
-            tables
-                .iter_external_def_id(tcx, krate)
-                .filter_map(|def_id| tables.to_static(tcx, def_id))
-                .collect()
-        }
+        tables.iter_def_ids(tcx, krate).filter_map(|def_id| tables.to_static(tcx, def_id)).collect()
     }
 
     fn foreign_module(
diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index c97c5ef8b4c8..f562213ac86c 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -93,11 +93,7 @@ impl<'tcx> Tables<'tcx> {
     }
 
     /// Iterate over the definitions of the given crate.
-    pub fn iter_external_def_id(
-        &self,
-        tcx: TyCtxt<'tcx>,
-        krate: CrateNum,
-    ) -> impl Iterator<Item = DefId> {
+    pub fn iter_def_ids(&self, tcx: TyCtxt<'tcx>, krate: CrateNum) -> impl Iterator<Item = DefId> {
         let num_definitions = tcx.num_def_ids(krate);
         (0..num_definitions)
             .map(move |i| DefId { krate, index: rustc_span::def_id::DefIndex::from_usize(i) })