Add method to convert internal to stable constructs

Cargo.lock

@@ -4524,6 +4524,7 @@ dependencies = [
  "rustc_middle",
  "rustc_span",
  "rustc_target",
+ "scoped-tls",
  "stable_mir",
  "tracing",
 ]

compiler/rustc_smir/Cargo.toml

@@ -9,6 +9,7 @@ rustc_hir = { path = "../rustc_hir" }
 rustc_middle = { path = "../rustc_middle" }
 rustc_span = { path = "../rustc_span" }
 rustc_target = { path = "../rustc_target" }
+scoped-tls = "1.0"
 stable_mir = {path = "../stable_mir" }
 tracing = "0.1"
 

compiler/rustc_smir/src/rustc_internal/mod.rs

@@ -3,21 +3,32 @@
 //! 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::Tables;
+use crate::rustc_smir::{Stable, Tables, TablesWrapper};
 use rustc_data_structures::fx;
 use rustc_data_structures::fx::FxIndexMap;
 use rustc_middle::mir::interpret::AllocId;
 use rustc_middle::ty;
 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::ty::IndexedVal;
+use std::cell::Cell;
+use std::cell::RefCell;
 use std::fmt::Debug;
 use std::hash::Hash;
 use std::ops::Index;
 
 mod internal;
 
+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 {
+    with_tables(|tables| item.internal(tables))
+}
+
 impl<'tcx> Index<stable_mir::DefId> for Tables<'tcx> {
     type Output = DefId;
 
@@ -125,18 +136,41 @@ pub fn crate_num(item: &stable_mir::Crate) -> CrateNum {
     item.id.into()
 }
 
+// A thread local variable that stores a pointer to the tables mapping between TyCtxt
+// datastructures and stable MIR datastructures
+scoped_thread_local! (static TLV: Cell<*const ()>);
+
+pub(crate) fn init<'tcx>(tables: &TablesWrapper<'tcx>, f: impl FnOnce()) {
+    assert!(!TLV.is_set());
+    let ptr: *const () = &tables as *const &_ as _;
+    TLV.set(&Cell::new(ptr), || {
+        f();
+    });
+}
+
+/// Loads the current context and calls a function with it.
+/// Do not nest these, as that will ICE.
+pub(crate) fn with_tables<'tcx, R>(f: impl FnOnce(&mut Tables<'tcx>) -> R) -> R {
+    assert!(TLV.is_set());
+    TLV.with(|tlv| {
+        let ptr = tlv.get();
+        assert!(!ptr.is_null());
+        let wrapper = unsafe { *(ptr as *const &TablesWrapper<'tcx>) };
+        let mut tables = wrapper.0.borrow_mut();
+        f(&mut *tables)
+    })
+}
+
 pub fn run(tcx: TyCtxt<'_>, f: impl FnOnce()) {
-    stable_mir::run(
-        Tables {
-            tcx,
-            def_ids: IndexMap::default(),
-            alloc_ids: IndexMap::default(),
-            spans: IndexMap::default(),
-            types: vec![],
-            instances: IndexMap::default(),
-        },
-        f,
-    );
+    let tables = TablesWrapper(RefCell::new(Tables {
+        tcx,
+        def_ids: IndexMap::default(),
+        alloc_ids: IndexMap::default(),
+        spans: IndexMap::default(),
+        types: vec![],
+        instances: IndexMap::default(),
+    }));
+    stable_mir::run(&tables, || init(&tables, f));
 }
 
 #[macro_export]
@@ -251,7 +285,7 @@ impl<K: PartialEq + Hash + Eq, V: Copy + Debug + PartialEq + IndexedVal> Index<V
 /// Trait used to translate a stable construct to its rustc counterpart.
 ///
 /// This is basically a mirror of [crate::rustc_smir::Stable].
-pub(crate) trait RustcInternal<'tcx> {
+pub trait RustcInternal<'tcx> {
     type T;
     fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T;
 }