Rollup merge of #132131 - celinval:smir-crate-defs, r=compiler-errors

[StableMIR] API to retrieve definitions from crates

Add functions to retrieve function definitions and static items from all crates (local and external).

For external crates, we're still missing items from trait implementation and primitives.

r? ````@compiler-errors:```` Do you know what is the best way to retrieve the associated items for primitives and trait implementations for external crates? Thanks!

compiler/rustc_smir/src/rustc_smir/context.rs

@@ -34,7 +34,7 @@ use stable_mir::{Crate, CrateDef, CrateItem, CrateNum, DefId, Error, Filename, I
 
 use crate::rustc_internal::RustcInternal;
 use crate::rustc_smir::builder::BodyBuilder;
-use crate::rustc_smir::{Stable, Tables, alloc, new_item_kind, smir_crate};
+use crate::rustc_smir::{Stable, Tables, alloc, filter_def_ids, new_item_kind, smir_crate};
 
 impl<'tcx> Context for TablesWrapper<'tcx> {
     fn target_info(&self) -> MachineInfo {
@@ -80,6 +80,20 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
             .collect()
     }
 
+    fn crate_functions(&self, crate_num: CrateNum) -> Vec<FnDef> {
+        let mut tables = self.0.borrow_mut();
+        let tcx = tables.tcx;
+        let krate = crate_num.internal(&mut *tables, tcx);
+        filter_def_ids(tcx, krate, |def_id| tables.to_fn_def(def_id))
+    }
+
+    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);
+        filter_def_ids(tcx, krate, |def_id| tables.to_static(def_id))
+    }
+
     fn foreign_module(
         &self,
         mod_def: stable_mir::ty::ForeignModuleDef,

compiler/rustc_smir/src/rustc_smir/mod.rs

@@ -15,8 +15,8 @@ use rustc_middle::mir::interpret::AllocId;
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
 use rustc_span::def_id::{CrateNum, DefId, LOCAL_CRATE};
 use stable_mir::abi::Layout;
-use stable_mir::mir::mono::InstanceDef;
-use stable_mir::ty::{MirConstId, Span, TyConstId};
+use stable_mir::mir::mono::{InstanceDef, StaticDef};
+use stable_mir::ty::{FnDef, MirConstId, Span, TyConstId};
 use stable_mir::{CtorKind, ItemKind};
 use tracing::debug;
 
@@ -79,6 +79,36 @@ impl<'tcx> Tables<'tcx> {
         };
         !must_override && self.tcx.is_mir_available(def_id)
     }
+
+    fn to_fn_def(&mut self, def_id: DefId) -> Option<FnDef> {
+        if matches!(self.tcx.def_kind(def_id), DefKind::Fn | DefKind::AssocFn) {
+            Some(self.fn_def(def_id))
+        } else {
+            None
+        }
+    }
+
+    fn to_static(&mut self, def_id: DefId) -> Option<StaticDef> {
+        matches!(self.tcx.def_kind(def_id), DefKind::Static { .. }).then(|| self.static_def(def_id))
+    }
+}
+
+/// Iterate over the definitions of the given crate.
+pub(crate) fn filter_def_ids<F, T>(tcx: TyCtxt<'_>, krate: CrateNum, mut func: F) -> Vec<T>
+where
+    F: FnMut(DefId) -> Option<T>,
+{
+    if krate == LOCAL_CRATE {
+        tcx.iter_local_def_id().filter_map(|did| func(did.to_def_id())).collect()
+    } else {
+        let num_definitions = tcx.num_extern_def_ids(krate);
+        (0..num_definitions)
+            .filter_map(move |i| {
+                let def_id = DefId { krate, index: rustc_span::def_id::DefIndex::from_usize(i) };
+                func(def_id)
+            })
+            .collect()
+    }
 }
 
 /// Build a stable mir crate from a given crate number.