4 files changed, 136 insertions, 63 deletions

diff --git a/compiler/rustc_smir/src/lib.rs b/compiler/rustc_smir/src/lib.rs
index 54d474db038..b00f0a1c153 100644
--- a/compiler/rustc_smir/src/lib.rs
+++ b/compiler/rustc_smir/src/lib.rs
@@ -11,6 +11,8 @@
     test(attr(allow(unused_variables), deny(warnings)))
 )]
 #![cfg_attr(not(feature = "default"), feature(rustc_private))]
+#![feature(local_key_cell_methods)]
+#![feature(ptr_metadata)]
 
 pub mod rustc_internal;
 pub mod stable_mir;
diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index 5998c8b6500..e643ad28c03 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -3,30 +3,49 @@
 //! For that, we define APIs that will temporarily be public to 3P that exposes rustc internal APIs
 //! until stable MIR is complete.
 
-use std::sync::RwLock;
-
-use crate::stable_mir;
+use crate::{
+    rustc_smir::Tables,
+    stable_mir::{self, with},
+};
+use rustc_middle::ty::TyCtxt;
 pub use rustc_span::def_id::{CrateNum, DefId};
 
-static DEF_ID_MAP: RwLock<Vec<DefId>> = RwLock::new(Vec::new());
+fn with_tables<R>(mut f: impl FnMut(&mut Tables<'_>) -> R) -> R {
+    let mut ret = None;
+    with(|tables| tables.rustc_tables(&mut |t| ret = Some(f(t))));
+    ret.unwrap()
+}
 
 pub fn item_def_id(item: &stable_mir::CrateItem) -> DefId {
-    DEF_ID_MAP.read().unwrap()[item.0]
+    with_tables(|t| t.item_def_id(item))
 }
 
 pub fn crate_item(did: DefId) -> stable_mir::CrateItem {
-    // FIXME: this becomes inefficient when we have too many ids
-    let mut map = DEF_ID_MAP.write().unwrap();
-    for (i, &d) in map.iter().enumerate() {
-        if d == did {
-            return stable_mir::CrateItem(i);
+    with_tables(|t| t.crate_item(did))
+}
+
+impl<'tcx> Tables<'tcx> {
+    pub fn item_def_id(&self, item: &stable_mir::CrateItem) -> DefId {
+        self.def_ids[item.0]
+    }
+
+    pub fn crate_item(&mut self, did: DefId) -> stable_mir::CrateItem {
+        // FIXME: this becomes inefficient when we have too many ids
+        for (i, &d) in self.def_ids.iter().enumerate() {
+            if d == did {
+                return stable_mir::CrateItem(i);
+            }
         }
+        let id = self.def_ids.len();
+        self.def_ids.push(did);
+        stable_mir::CrateItem(id)
     }
-    let id = map.len();
-    map.push(did);
-    stable_mir::CrateItem(id)
 }
 
 pub fn crate_num(item: &stable_mir::Crate) -> CrateNum {
     item.id.into()
 }
+
+pub fn run(tcx: TyCtxt<'_>, f: impl FnOnce()) {
+    crate::stable_mir::run(Tables { tcx, def_ids: vec![] }, f);
+}
diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index 241cd182059..81ecce415d5 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -7,55 +7,36 @@
 //!
 //! For now, we are developing everything inside `rustc`, thus, we keep this module private.
 
-use crate::{
-    rustc_internal::{crate_item, item_def_id},
-    stable_mir::{self},
-};
-use rustc_middle::ty::{tls::with, TyCtxt};
-use rustc_span::def_id::{CrateNum, LOCAL_CRATE};
+use crate::stable_mir::{self, Context};
+use rustc_middle::ty::TyCtxt;
+use rustc_span::def_id::{CrateNum, DefId, LOCAL_CRATE};
 use tracing::debug;
 
-/// Get information about the local crate.
-pub fn local_crate() -> stable_mir::Crate {
-    with(|tcx| smir_crate(tcx, LOCAL_CRATE))
-}
+impl<'tcx> Context for Tables<'tcx> {
+    fn local_crate(&self) -> stable_mir::Crate {
+        smir_crate(self.tcx, LOCAL_CRATE)
+    }
 
-/// Retrieve a list of all external crates.
-pub fn external_crates() -> Vec<stable_mir::Crate> {
-    with(|tcx| tcx.crates(()).iter().map(|crate_num| smir_crate(tcx, *crate_num)).collect())
-}
+    fn external_crates(&self) -> Vec<stable_mir::Crate> {
+        self.tcx.crates(()).iter().map(|crate_num| smir_crate(self.tcx, *crate_num)).collect()
+    }
 
-/// Find a crate with the given name.
-pub fn find_crate(name: &str) -> Option<stable_mir::Crate> {
-    with(|tcx| {
-        [LOCAL_CRATE].iter().chain(tcx.crates(()).iter()).find_map(|crate_num| {
-            let crate_name = tcx.crate_name(*crate_num).to_string();
-            (name == crate_name).then(|| smir_crate(tcx, *crate_num))
+    fn find_crate(&self, name: &str) -> Option<stable_mir::Crate> {
+        [LOCAL_CRATE].iter().chain(self.tcx.crates(()).iter()).find_map(|crate_num| {
+            let crate_name = self.tcx.crate_name(*crate_num).to_string();
+            (name == crate_name).then(|| smir_crate(self.tcx, *crate_num))
         })
-    })
-}
-
-/// Retrieve all items of the local crate that have a MIR associated with them.
-pub fn all_local_items() -> stable_mir::CrateItems {
-    with(|tcx| tcx.mir_keys(()).iter().map(|item| crate_item(item.to_def_id())).collect())
-}
-
-pub fn entry_fn() -> Option<stable_mir::CrateItem> {
-    with(|tcx| Some(crate_item(tcx.entry_fn(())?.0)))
-}
-
-/// Build a stable mir crate from a given crate number.
-fn smir_crate(tcx: TyCtxt<'_>, crate_num: CrateNum) -> stable_mir::Crate {
-    let crate_name = tcx.crate_name(crate_num).to_string();
-    let is_local = crate_num == LOCAL_CRATE;
-    debug!(?crate_name, ?crate_num, "smir_crate");
-    stable_mir::Crate { id: crate_num.into(), name: crate_name, is_local }
-}
+    }
 
-pub fn mir_body(item: &stable_mir::CrateItem) -> stable_mir::mir::Body {
-    with(|tcx| {
-        let def_id = item_def_id(item);
-        let mir = tcx.optimized_mir(def_id);
+    fn all_local_items(&mut self) -> stable_mir::CrateItems {
+        self.tcx.mir_keys(()).iter().map(|item| self.crate_item(item.to_def_id())).collect()
+    }
+    fn entry_fn(&mut self) -> Option<stable_mir::CrateItem> {
+        Some(self.crate_item(self.tcx.entry_fn(())?.0))
+    }
+    fn mir_body(&self, item: &stable_mir::CrateItem) -> stable_mir::mir::Body {
+        let def_id = self.item_def_id(item);
+        let mir = self.tcx.optimized_mir(def_id);
         stable_mir::mir::Body {
             blocks: mir
                 .basic_blocks
@@ -66,7 +47,24 @@ pub fn mir_body(item: &stable_mir::CrateItem) -> stable_mir::mir::Body {
                 })
                 .collect(),
         }
-    })
+    }
+
+    fn rustc_tables(&mut self, f: &mut dyn FnMut(&mut Tables<'_>)) {
+        f(self)
+    }
+}
+
+pub struct Tables<'tcx> {
+    pub tcx: TyCtxt<'tcx>,
+    pub def_ids: Vec<DefId>,
+}
+
+/// Build a stable mir crate from a given crate number.
+fn smir_crate(tcx: TyCtxt<'_>, crate_num: CrateNum) -> stable_mir::Crate {
+    let crate_name = tcx.crate_name(crate_num).to_string();
+    let is_local = crate_num == LOCAL_CRATE;
+    debug!(?crate_name, ?crate_num, "smir_crate");
+    stable_mir::Crate { id: crate_num.into(), name: crate_name, is_local }
 }
 
 fn rustc_statement_to_statement(
diff --git a/compiler/rustc_smir/src/stable_mir/mod.rs b/compiler/rustc_smir/src/stable_mir/mod.rs
index 1d2efb5eab9..612777b9c75 100644
--- a/compiler/rustc_smir/src/stable_mir/mod.rs
+++ b/compiler/rustc_smir/src/stable_mir/mod.rs
@@ -11,7 +11,14 @@
 //! There shouldn't be any direct references to internal compiler constructs in this module.
 //! If you need an internal construct, consider using `rustc_internal` or `rustc_smir`.
 
+use std::cell::Cell;
+
+use crate::rustc_smir::Tables;
+
+use self::ty::{Ty, TyKind};
+
 pub mod mir;
+pub mod ty;
 
 /// Use String for now but we should replace it.
 pub type Symbol = String;
@@ -41,7 +48,7 @@ pub struct CrateItem(pub(crate) DefId);
 
 impl CrateItem {
     pub fn body(&self) -> mir::Body {
-        crate::rustc_smir::mir_body(self)
+        with(|cx| cx.mir_body(self))
     }
 }
 
@@ -49,25 +56,72 @@ impl CrateItem {
 /// crate defines that. This is usually `main`, but could be
 /// `start` if the crate is a no-std crate.
 pub fn entry_fn() -> Option<CrateItem> {
-    crate::rustc_smir::entry_fn()
+    with(|cx| cx.entry_fn())
 }
 
 /// Access to the local crate.
 pub fn local_crate() -> Crate {
-    crate::rustc_smir::local_crate()
+    with(|cx| cx.local_crate())
 }
 
 /// Try to find a crate with the given name.
 pub fn find_crate(name: &str) -> Option<Crate> {
-    crate::rustc_smir::find_crate(name)
+    with(|cx| cx.find_crate(name))
 }
 
 /// Try to find a crate with the given name.
 pub fn external_crates() -> Vec<Crate> {
-    crate::rustc_smir::external_crates()
+    with(|cx| cx.external_crates())
 }
 
 /// Retrieve all items in the local crate that have a MIR associated with them.
 pub fn all_local_items() -> CrateItems {
-    crate::rustc_smir::all_local_items()
+    with(|cx| cx.all_local_items())
+}
+
+pub trait Context {
+    fn entry_fn(&mut self) -> Option<CrateItem>;
+    /// Retrieve all items of the local crate that have a MIR associated with them.
+    fn all_local_items(&mut self) -> CrateItems;
+    fn mir_body(&mut self, item: &CrateItem) -> mir::Body;
+    /// Get information about the local crate.
+    fn local_crate(&self) -> Crate;
+    /// Retrieve a list of all external crates.
+    fn external_crates(&self) -> Vec<Crate>;
+
+    /// Find a crate with the given name.
+    fn find_crate(&self, name: &str) -> Option<Crate>;
+
+    /// Obtain the representation of a type.
+    fn ty_kind(&mut self, ty: Ty) -> TyKind;
+
+    /// HACK: Until we have fully stable consumers, we need an escape hatch
+    /// to get `DefId`s out of `CrateItem`s.
+    fn rustc_tables(&mut self, f: &mut dyn FnMut(&mut Tables<'_>));
+}
+
+thread_local! {
+    /// A thread local variable that stores a pointer to the tables mapping between TyCtxt
+    /// datastructures and stable MIR datastructures.
+    static TLV: Cell<*mut ()> = const { Cell::new(std::ptr::null_mut()) };
+}
+
+pub fn run(mut context: impl Context, f: impl FnOnce()) {
+    assert!(TLV.get().is_null());
+    fn g<'a>(mut context: &mut (dyn Context + 'a), f: impl FnOnce()) {
+        TLV.set(&mut context as *mut &mut _ as _);
+        f();
+        TLV.replace(std::ptr::null_mut());
+    }
+    g(&mut context, f);
+}
+
+/// Loads the current context and calls a function with it.
+/// Do not nest these, as that will ICE.
+pub(crate) fn with<R>(f: impl FnOnce(&mut dyn Context) -> R) -> R {
+    let ptr = TLV.replace(std::ptr::null_mut()) as *mut &mut dyn Context;
+    assert!(!ptr.is_null());
+    let ret = f(unsafe { *ptr });
+    TLV.set(ptr as _);
+    ret
 }