1 files changed, 28 insertions, 11 deletions

diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index 4d2a51822..4bac98909 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -3,7 +3,7 @@
 //! For that, we define APIs that will temporarily be public to 3P that exposes rustc internal APIs
 //! until stable MIR is complete.
 
-use crate::rustc_smir::{Stable, Tables, TablesWrapper};
+use crate::rustc_smir::{context::TablesWrapper, Stable, Tables};
 use rustc_data_structures::fx;
 use rustc_data_structures::fx::FxIndexMap;
 use rustc_middle::mir::interpret::AllocId;
@@ -12,7 +12,9 @@ use rustc_middle::ty::TyCtxt;
 use rustc_span::def_id::{CrateNum, DefId};
 use rustc_span::Span;
 use scoped_tls::scoped_thread_local;
+use stable_mir::abi::Layout;
 use stable_mir::ty::IndexedVal;
+use stable_mir::Error;
 use std::cell::Cell;
 use std::cell::RefCell;
 use std::fmt::Debug;
@@ -20,12 +22,13 @@ use std::hash::Hash;
 use std::ops::Index;
 
 mod internal;
+pub mod pretty;
 
-pub fn stable<'tcx, S: Stable<'tcx>>(item: &S) -> S::T {
+pub fn stable<'tcx, S: Stable<'tcx>>(item: S) -> S::T {
     with_tables(|tables| item.stable(tables))
 }
 
-pub fn internal<'tcx, S: RustcInternal<'tcx>>(item: &S) -> S::T {
+pub fn internal<'tcx, S: RustcInternal<'tcx>>(item: S) -> S::T {
     with_tables(|tables| item.internal(tables))
 }
 
@@ -104,6 +107,10 @@ impl<'tcx> Tables<'tcx> {
         stable_mir::ty::RegionDef(self.create_def_id(did))
     }
 
+    pub fn coroutine_witness_def(&mut self, did: DefId) -> stable_mir::ty::CoroutineWitnessDef {
+        stable_mir::ty::CoroutineWitnessDef(self.create_def_id(did))
+    }
+
     pub fn prov(&mut self, aid: AllocId) -> stable_mir::ty::Prov {
         stable_mir::ty::Prov(self.create_alloc_id(aid))
     }
@@ -112,7 +119,7 @@ impl<'tcx> Tables<'tcx> {
         self.def_ids.create_or_fetch(did)
     }
 
-    fn create_alloc_id(&mut self, aid: AllocId) -> stable_mir::AllocId {
+    pub(crate) fn create_alloc_id(&mut self, aid: AllocId) -> stable_mir::mir::alloc::AllocId {
         self.alloc_ids.create_or_fetch(aid)
     }
 
@@ -130,6 +137,10 @@ impl<'tcx> Tables<'tcx> {
     pub(crate) fn static_def(&mut self, did: DefId) -> stable_mir::mir::mono::StaticDef {
         stable_mir::mir::mono::StaticDef(self.create_def_id(did))
     }
+
+    pub(crate) fn layout_id(&mut self, layout: rustc_target::abi::Layout<'tcx>) -> Layout {
+        self.layouts.create_or_fetch(layout)
+    }
 }
 
 pub fn crate_num(item: &stable_mir::Crate) -> CrateNum {
@@ -140,12 +151,13 @@ pub fn crate_num(item: &stable_mir::Crate) -> CrateNum {
 // datastructures and stable MIR datastructures
 scoped_thread_local! (static TLV: Cell<*const ()>);
 
-pub(crate) fn init<'tcx>(tables: &TablesWrapper<'tcx>, f: impl FnOnce()) {
+pub(crate) fn init<'tcx, F, T>(tables: &TablesWrapper<'tcx>, f: F) -> T
+where
+    F: FnOnce() -> T,
+{
     assert!(!TLV.is_set());
     let ptr = tables as *const _ as *const ();
-    TLV.set(&Cell::new(ptr), || {
-        f();
-    });
+    TLV.set(&Cell::new(ptr), || f())
 }
 
 /// Loads the current context and calls a function with it.
@@ -161,7 +173,10 @@ pub(crate) fn with_tables<'tcx, R>(f: impl FnOnce(&mut Tables<'tcx>) -> R) -> R
     })
 }
 
-pub fn run(tcx: TyCtxt<'_>, f: impl FnOnce()) {
+pub fn run<F, T>(tcx: TyCtxt<'_>, f: F) -> Result<T, Error>
+where
+    F: FnOnce() -> T,
+{
     let tables = TablesWrapper(RefCell::new(Tables {
         tcx,
         def_ids: IndexMap::default(),
@@ -170,8 +185,9 @@ pub fn run(tcx: TyCtxt<'_>, f: impl FnOnce()) {
         types: IndexMap::default(),
         instances: IndexMap::default(),
         constants: IndexMap::default(),
+        layouts: IndexMap::default(),
     }));
-    stable_mir::run(&tables, || init(&tables, f));
+    stable_mir::compiler_interface::run(&tables, || init(&tables, f))
 }
 
 #[macro_export]
@@ -237,7 +253,8 @@ macro_rules! run {
                 queries.global_ctxt().unwrap().enter(|tcx| {
                     rustc_internal::run(tcx, || {
                         self.result = Some((self.callback)(tcx));
-                    });
+                    })
+                    .unwrap();
                     if self.result.as_ref().is_some_and(|val| val.is_continue()) {
                         Compilation::Continue
                     } else {