Rollup merge of #117010 - celinval:smir-internal, r=oli-obk

Add method to convert internal to stable constructs

This is an alternative implementation to #116999. I believe we can still improve the logic a bit here, but I wanted to see which direction we should go first.

In this implementation, the API is simpler and we keep Tables somewhat private. The definition is still public though, since we have to expose the Stable trait. However, there's a cost of keeping another thread-local and using `Rc`, but I'm hoping it will be a small cost.

r? ``@oli-obk``
r? ``@spastorino``

Cargo.lock

@@ -4525,6 +4525,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 = tables as *const _ as *const ();
+    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 = ptr as *const TablesWrapper<'tcx>;
+        let mut tables = unsafe { (*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;
 }