Change how smir iterate over all external crate defs

Instead of trying to gather all definitions by walking the exported
modules, types and traits, add a query to expose the method to
retrieve the number of defs in a crate. Then, generate the `def_ids`
using this value.

For the local crate, we still use `hir_crate_items` query that provides
an iterator to all local definitions.

Author: celinval

compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs

@@ -387,6 +387,7 @@ provide! { tcx, def_id, other, cdata,
     crate_hash => { cdata.root.header.hash }
     crate_host_hash => { cdata.host_hash }
     crate_name => { cdata.root.header.name }
+    num_def_ids => { cdata.num_def_ids() }
 
     extra_filename => { cdata.root.extra_filename.clone() }
 

compiler/rustc_middle/src/hir/mod.rs

@@ -233,4 +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();
 }

compiler/rustc_middle/src/query/mod.rs

@@ -1812,6 +1812,16 @@ rustc_queries! {
         desc { |tcx| "computing crate imported by `{}`", tcx.def_path_str(def_id) }
     }
 
+    /// 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
+    }
+
     query lib_features(_: CrateNum) -> &'tcx LibFeatures {
         desc { "calculating the lib features defined in a crate" }
         separate_provide_extern

compiler/rustc_smir/src/rustc_smir/context.rs

@@ -90,7 +90,10 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
                 .filter_map(|local_id| tables.to_fn_def(tcx, local_id.to_def_id()))
                 .collect()
         } else {
-            tables.filter_map_items(tcx, krate, |tables, tcx, def_id| tables.to_fn_def(tcx, def_id))
+            tables
+                .iter_external_def_id(tcx, krate)
+                .filter_map(|def_id| tables.to_fn_def(tcx, def_id))
+                .collect()
         }
     }
 
@@ -104,7 +107,10 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
                 .filter_map(|local_id| tables.to_static(tcx, local_id.to_def_id()))
                 .collect()
         } else {
-            tables.filter_map_items(tcx, krate, |tables, tcx, def_id| tables.to_static(tcx, def_id))
+            tables
+                .iter_external_def_id(tcx, krate)
+                .filter_map(|def_id| tables.to_static(tcx, def_id))
+                .collect()
         }
     }
 

compiler/rustc_smir/src/rustc_smir/mod.rs

@@ -12,7 +12,6 @@ use std::ops::RangeInclusive;
 use rustc_hir::def::DefKind;
 use rustc_middle::mir;
 use rustc_middle::mir::interpret::AllocId;
-use rustc_middle::ty::data_structures::HashSet;
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
 use rustc_span::def_id::{CrateNum, DefId, LOCAL_CRATE};
 use stable_mir::abi::Layout;
@@ -93,34 +92,15 @@ impl<'tcx> Tables<'tcx> {
         matches!(tcx.def_kind(def_id), DefKind::Static { .. }).then(|| self.static_def(def_id))
     }
 
-    fn filter_map_items<T, F>(&mut self, tcx: TyCtxt<'tcx>, krate: CrateNum, func: F) -> Vec<T>
-    where
-        F: Fn(&mut Tables<'tcx>, TyCtxt<'tcx>, DefId) -> Option<T>,
-    {
-        let crate_def = krate.as_def_id();
-        let mut queue = vec![crate_def];
-        queue.extend_from_slice(tcx.trait_impls_in_crate(krate));
-        let mut visited = HashSet::default();
-        let mut result = vec![];
-        while let Some(def_id) = queue.pop() {
-            if visited.insert(def_id) {
-                result.extend(func(self, tcx, def_id));
-                match tcx.def_kind(def_id) {
-                    DefKind::Mod | DefKind::ForeignMod | DefKind::Trait => queue.extend(
-                        tcx.module_children(def_id).iter().filter_map(|item| item.res.opt_def_id()),
-                    ),
-                    DefKind::Impl { .. } => queue.extend(
-                        tcx.associated_items(def_id).in_definition_order().map(|item| item.def_id),
-                    ),
-                    DefKind::Struct | DefKind::Enum | DefKind::Union => {
-                        queue.extend(tcx.inherent_impls(def_id));
-                    }
-                    _ => {}
-                }
-            }
-        }
-
-        result
+    /// Iterate over the definitions of the given crate.
+    pub fn iter_external_def_id(
+        &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) })
     }
 }
 

tests/ui-fulldeps/stable-mir/check_crate_defs.rs

@@ -53,6 +53,7 @@ fn test_stable_mir() -> ControlFlow<()> {
             "std::option::Option::<T>::is_some",
             "std::ptr::swap",
             "<std::slice::Iter<'a, T> as std::iter::Iterator>::next",
+            "core::num::<impl u8>::abs_diff",
         ],
     );
     // Ensure nothing crashes. There is no public static in core that we can test here.