6 files changed, 417 insertions, 77 deletions

diff --git a/compiler/rustc_smir/src/lib.rs b/compiler/rustc_smir/src/lib.rs
index 54d474db0..b00f0a1c1 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 5998c8b65..609a04d26 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![], types: vec![] }, f);
+}
diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index 4dad3c6bc..5572108f4 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -7,41 +7,107 @@
 //!
 //! 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, ty::TyKind, Context};
+use rustc_middle::ty::{self, 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))
         })
-    })
+    }
+
+    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(&mut 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
+                .iter()
+                .map(|block| stable_mir::mir::BasicBlock {
+                    terminator: rustc_terminator_to_terminator(block.terminator()),
+                    statements: block.statements.iter().map(rustc_statement_to_statement).collect(),
+                })
+                .collect(),
+            locals: mir.local_decls.iter().map(|decl| self.intern_ty(decl.ty)).collect(),
+        }
+    }
+
+    fn rustc_tables(&mut self, f: &mut dyn FnMut(&mut Tables<'_>)) {
+        f(self)
+    }
+
+    fn ty_kind(&mut self, ty: crate::stable_mir::ty::Ty) -> TyKind {
+        self.rustc_ty_to_ty(self.types[ty.0])
+    }
 }
 
-/// 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 struct Tables<'tcx> {
+    pub tcx: TyCtxt<'tcx>,
+    pub def_ids: Vec<DefId>,
+    pub types: Vec<Ty<'tcx>>,
 }
 
-pub fn entry_fn() -> Option<stable_mir::CrateItem> {
-    with(|tcx| Some(crate_item(tcx.entry_fn(())?.0)))
+impl<'tcx> Tables<'tcx> {
+    fn rustc_ty_to_ty(&mut self, ty: Ty<'tcx>) -> TyKind {
+        match ty.kind() {
+            ty::Bool => TyKind::Bool,
+            ty::Char => todo!(),
+            ty::Int(_) => todo!(),
+            ty::Uint(_) => todo!(),
+            ty::Float(_) => todo!(),
+            ty::Adt(_, _) => todo!(),
+            ty::Foreign(_) => todo!(),
+            ty::Str => todo!(),
+            ty::Array(_, _) => todo!(),
+            ty::Slice(_) => todo!(),
+            ty::RawPtr(_) => todo!(),
+            ty::Ref(_, _, _) => todo!(),
+            ty::FnDef(_, _) => todo!(),
+            ty::FnPtr(_) => todo!(),
+            ty::Placeholder(..) => todo!(),
+            ty::Dynamic(_, _, _) => todo!(),
+            ty::Closure(_, _) => todo!(),
+            ty::Generator(_, _, _) => todo!(),
+            ty::GeneratorWitness(_) => todo!(),
+            ty::GeneratorWitnessMIR(_, _) => todo!(),
+            ty::Never => todo!(),
+            ty::Tuple(fields) => {
+                TyKind::Tuple(fields.iter().map(|ty| self.intern_ty(ty)).collect())
+            }
+            ty::Alias(_, _) => todo!(),
+            ty::Param(_) => todo!(),
+            ty::Bound(_, _) => todo!(),
+            ty::Infer(_) => todo!(),
+            ty::Error(_) => todo!(),
+        }
+    }
+
+    fn intern_ty(&mut self, ty: Ty<'tcx>) -> stable_mir::ty::Ty {
+        if let Some(id) = self.types.iter().position(|&t| t == ty) {
+            return stable_mir::ty::Ty(id);
+        }
+        let id = self.types.len();
+        self.types.push(ty);
+        stable_mir::ty::Ty(id)
+    }
 }
 
 /// Build a stable mir crate from a given crate number.
@@ -52,23 +118,6 @@ fn smir_crate(tcx: TyCtxt<'_>, crate_num: CrateNum) -> stable_mir::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);
-        stable_mir::mir::Body {
-            blocks: mir
-                .basic_blocks
-                .iter()
-                .map(|block| stable_mir::mir::BasicBlock {
-                    terminator: rustc_terminator_to_terminator(block.terminator()),
-                    statements: block.statements.iter().map(rustc_statement_to_statement).collect(),
-                })
-                .collect(),
-        }
-    })
-}
-
 fn rustc_statement_to_statement(
     s: &rustc_middle::mir::Statement<'_>,
 ) -> stable_mir::mir::Statement {
@@ -93,10 +142,10 @@ fn rustc_statement_to_statement(
     }
 }
 
-fn rustc_rvalue_to_rvalue(rvalue: &rustc_middle::mir::Rvalue<'_>) -> stable_mir::mir::Operand {
+fn rustc_rvalue_to_rvalue(rvalue: &rustc_middle::mir::Rvalue<'_>) -> stable_mir::mir::Rvalue {
     use rustc_middle::mir::Rvalue::*;
     match rvalue {
-        Use(op) => rustc_op_to_op(op),
+        Use(op) => stable_mir::mir::Rvalue::Use(rustc_op_to_op(op)),
         Repeat(_, _) => todo!(),
         Ref(_, _, _) => todo!(),
         ThreadLocalRef(_) => todo!(),
@@ -104,9 +153,15 @@ fn rustc_rvalue_to_rvalue(rvalue: &rustc_middle::mir::Rvalue<'_>) -> stable_mir:
         Len(_) => todo!(),
         Cast(_, _, _) => todo!(),
         BinaryOp(_, _) => todo!(),
-        CheckedBinaryOp(_, _) => todo!(),
+        CheckedBinaryOp(bin_op, ops) => stable_mir::mir::Rvalue::CheckedBinaryOp(
+            rustc_bin_op_to_bin_op(bin_op),
+            rustc_op_to_op(&ops.0),
+            rustc_op_to_op(&ops.1),
+        ),
         NullaryOp(_, _) => todo!(),
-        UnaryOp(_, _) => todo!(),
+        UnaryOp(un_op, op) => {
+            stable_mir::mir::Rvalue::UnaryOp(rustc_un_op_to_un_op(un_op), rustc_op_to_op(op))
+        }
         Discriminant(_) => todo!(),
         Aggregate(_, _) => todo!(),
         ShallowInitBox(_, _) => todo!(),
@@ -124,8 +179,112 @@ fn rustc_op_to_op(op: &rustc_middle::mir::Operand<'_>) -> stable_mir::mir::Opera
 }
 
 fn rustc_place_to_place(place: &rustc_middle::mir::Place<'_>) -> stable_mir::mir::Place {
-    assert_eq!(&place.projection[..], &[]);
-    stable_mir::mir::Place { local: place.local.as_usize() }
+    stable_mir::mir::Place {
+        local: place.local.as_usize(),
+        projection: format!("{:?}", place.projection),
+    }
+}
+
+fn rustc_unwind_to_unwind(
+    unwind: &rustc_middle::mir::UnwindAction,
+) -> stable_mir::mir::UnwindAction {
+    use rustc_middle::mir::UnwindAction;
+    match unwind {
+        UnwindAction::Continue => stable_mir::mir::UnwindAction::Continue,
+        UnwindAction::Unreachable => stable_mir::mir::UnwindAction::Unreachable,
+        UnwindAction::Terminate => stable_mir::mir::UnwindAction::Terminate,
+        UnwindAction::Cleanup(bb) => stable_mir::mir::UnwindAction::Cleanup(bb.as_usize()),
+    }
+}
+
+fn rustc_assert_msg_to_msg<'tcx>(
+    assert_message: &rustc_middle::mir::AssertMessage<'tcx>,
+) -> stable_mir::mir::AssertMessage {
+    use rustc_middle::mir::AssertKind;
+    match assert_message {
+        AssertKind::BoundsCheck { len, index } => stable_mir::mir::AssertMessage::BoundsCheck {
+            len: rustc_op_to_op(len),
+            index: rustc_op_to_op(index),
+        },
+        AssertKind::Overflow(bin_op, op1, op2) => stable_mir::mir::AssertMessage::Overflow(
+            rustc_bin_op_to_bin_op(bin_op),
+            rustc_op_to_op(op1),
+            rustc_op_to_op(op2),
+        ),
+        AssertKind::OverflowNeg(op) => {
+            stable_mir::mir::AssertMessage::OverflowNeg(rustc_op_to_op(op))
+        }
+        AssertKind::DivisionByZero(op) => {
+            stable_mir::mir::AssertMessage::DivisionByZero(rustc_op_to_op(op))
+        }
+        AssertKind::RemainderByZero(op) => {
+            stable_mir::mir::AssertMessage::RemainderByZero(rustc_op_to_op(op))
+        }
+        AssertKind::ResumedAfterReturn(generator) => {
+            stable_mir::mir::AssertMessage::ResumedAfterReturn(rustc_generator_to_generator(
+                generator,
+            ))
+        }
+        AssertKind::ResumedAfterPanic(generator) => {
+            stable_mir::mir::AssertMessage::ResumedAfterPanic(rustc_generator_to_generator(
+                generator,
+            ))
+        }
+        AssertKind::MisalignedPointerDereference { required, found } => {
+            stable_mir::mir::AssertMessage::MisalignedPointerDereference {
+                required: rustc_op_to_op(required),
+                found: rustc_op_to_op(found),
+            }
+        }
+    }
+}
+
+fn rustc_bin_op_to_bin_op(bin_op: &rustc_middle::mir::BinOp) -> stable_mir::mir::BinOp {
+    use rustc_middle::mir::BinOp;
+    match bin_op {
+        BinOp::Add => stable_mir::mir::BinOp::Add,
+        BinOp::Sub => stable_mir::mir::BinOp::Sub,
+        BinOp::Mul => stable_mir::mir::BinOp::Mul,
+        BinOp::Div => stable_mir::mir::BinOp::Div,
+        BinOp::Rem => stable_mir::mir::BinOp::Rem,
+        BinOp::BitXor => stable_mir::mir::BinOp::BitXor,
+        BinOp::BitAnd => stable_mir::mir::BinOp::BitAnd,
+        BinOp::BitOr => stable_mir::mir::BinOp::BitOr,
+        BinOp::Shl => stable_mir::mir::BinOp::Shl,
+        BinOp::Shr => stable_mir::mir::BinOp::Shr,
+        BinOp::Eq => stable_mir::mir::BinOp::Eq,
+        BinOp::Lt => stable_mir::mir::BinOp::Lt,
+        BinOp::Le => stable_mir::mir::BinOp::Le,
+        BinOp::Ne => stable_mir::mir::BinOp::Ne,
+        BinOp::Ge => stable_mir::mir::BinOp::Ge,
+        BinOp::Gt => stable_mir::mir::BinOp::Gt,
+        BinOp::Offset => stable_mir::mir::BinOp::Offset,
+    }
+}
+
+fn rustc_un_op_to_un_op(unary_op: &rustc_middle::mir::UnOp) -> stable_mir::mir::UnOp {
+    use rustc_middle::mir::UnOp;
+    match unary_op {
+        UnOp::Not => stable_mir::mir::UnOp::Not,
+        UnOp::Neg => stable_mir::mir::UnOp::Neg,
+    }
+}
+
+fn rustc_generator_to_generator(
+    generator: &rustc_hir::GeneratorKind,
+) -> stable_mir::mir::GeneratorKind {
+    use rustc_hir::{AsyncGeneratorKind, GeneratorKind};
+    match generator {
+        GeneratorKind::Async(async_gen) => {
+            let async_gen = match async_gen {
+                AsyncGeneratorKind::Block => stable_mir::mir::AsyncGeneratorKind::Block,
+                AsyncGeneratorKind::Closure => stable_mir::mir::AsyncGeneratorKind::Closure,
+                AsyncGeneratorKind::Fn => stable_mir::mir::AsyncGeneratorKind::Fn,
+            };
+            stable_mir::mir::GeneratorKind::Async(async_gen)
+        }
+        GeneratorKind::Gen => stable_mir::mir::GeneratorKind::Gen,
+    }
 }
 
 fn rustc_terminator_to_terminator(
@@ -150,11 +309,29 @@ fn rustc_terminator_to_terminator(
         Terminate => Terminator::Abort,
         Return => Terminator::Return,
         Unreachable => Terminator::Unreachable,
-        Drop { .. } => todo!(),
-        Call { .. } => todo!(),
-        Assert { .. } => todo!(),
+        Drop { place, target, unwind, replace: _ } => Terminator::Drop {
+            place: rustc_place_to_place(place),
+            target: target.as_usize(),
+            unwind: rustc_unwind_to_unwind(unwind),
+        },
+        Call { func, args, destination, target, unwind, from_hir_call: _, fn_span: _ } => {
+            Terminator::Call {
+                func: rustc_op_to_op(func),
+                args: args.iter().map(|arg| rustc_op_to_op(arg)).collect(),
+                destination: rustc_place_to_place(destination),
+                target: target.map(|t| t.as_usize()),
+                unwind: rustc_unwind_to_unwind(unwind),
+            }
+        }
+        Assert { cond, expected, msg, target, unwind } => Terminator::Assert {
+            cond: rustc_op_to_op(cond),
+            expected: *expected,
+            msg: rustc_assert_msg_to_msg(msg),
+            target: target.as_usize(),
+            unwind: rustc_unwind_to_unwind(unwind),
+        },
         Yield { .. } => todo!(),
-        GeneratorDrop => todo!(),
+        GeneratorDrop => Terminator::GeneratorDrop,
         FalseEdge { .. } => todo!(),
         FalseUnwind { .. } => todo!(),
         InlineAsm { .. } => todo!(),
diff --git a/compiler/rustc_smir/src/stable_mir/mir/body.rs b/compiler/rustc_smir/src/stable_mir/mir/body.rs
index c504065c9..6328c35aa 100644
--- a/compiler/rustc_smir/src/stable_mir/mir/body.rs
+++ b/compiler/rustc_smir/src/stable_mir/mir/body.rs
@@ -1,6 +1,9 @@
+use crate::stable_mir::ty::Ty;
+
 #[derive(Clone, Debug)]
 pub struct Body {
     pub blocks: Vec<BasicBlock>,
+    pub locals: Vec<Ty>,
 }
 
 #[derive(Clone, Debug)]
@@ -26,30 +29,99 @@ pub enum Terminator {
     Drop {
         place: Place,
         target: usize,
-        unwind: Option<usize>,
+        unwind: UnwindAction,
     },
     Call {
         func: Operand,
         args: Vec<Operand>,
         destination: Place,
         target: Option<usize>,
-        cleanup: Option<usize>,
+        unwind: UnwindAction,
     },
     Assert {
         cond: Operand,
         expected: bool,
-        msg: String,
+        msg: AssertMessage,
         target: usize,
-        cleanup: Option<usize>,
+        unwind: UnwindAction,
     },
+    GeneratorDrop,
+}
+
+#[derive(Clone, Debug)]
+pub enum UnwindAction {
+    Continue,
+    Unreachable,
+    Terminate,
+    Cleanup(usize),
+}
+
+#[derive(Clone, Debug)]
+pub enum AssertMessage {
+    BoundsCheck { len: Operand, index: Operand },
+    Overflow(BinOp, Operand, Operand),
+    OverflowNeg(Operand),
+    DivisionByZero(Operand),
+    RemainderByZero(Operand),
+    ResumedAfterReturn(GeneratorKind),
+    ResumedAfterPanic(GeneratorKind),
+    MisalignedPointerDereference { required: Operand, found: Operand },
+}
+
+#[derive(Clone, Debug)]
+pub enum BinOp {
+    Add,
+    Sub,
+    Mul,
+    Div,
+    Rem,
+    BitXor,
+    BitAnd,
+    BitOr,
+    Shl,
+    Shr,
+    Eq,
+    Lt,
+    Le,
+    Ne,
+    Ge,
+    Gt,
+    Offset,
+}
+
+#[derive(Clone, Debug)]
+pub enum UnOp {
+    Not,
+    Neg,
+}
+
+#[derive(Clone, Debug)]
+pub enum GeneratorKind {
+    Async(AsyncGeneratorKind),
+    Gen,
+}
+
+#[derive(Clone, Debug)]
+pub enum AsyncGeneratorKind {
+    Block,
+    Closure,
+    Fn,
 }
 
 #[derive(Clone, Debug)]
 pub enum Statement {
-    Assign(Place, Operand),
+    Assign(Place, Rvalue),
     Nop,
 }
 
+// FIXME this is incomplete
+#[derive(Clone, Debug)]
+pub enum Rvalue {
+    Use(Operand),
+    CheckedBinaryOp(BinOp, Operand, Operand),
+    UnaryOp(UnOp, Operand),
+}
+
 #[derive(Clone, Debug)]
 pub enum Operand {
     Copy(Place),
@@ -60,6 +132,7 @@ pub enum Operand {
 #[derive(Clone, Debug)]
 pub struct Place {
     pub local: usize,
+    pub projection: String,
 }
 
 #[derive(Clone, Debug)]
diff --git a/compiler/rustc_smir/src/stable_mir/mod.rs b/compiler/rustc_smir/src/stable_mir/mod.rs
index 1d2efb5ea..612777b9c 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
 }
diff --git a/compiler/rustc_smir/src/stable_mir/ty.rs b/compiler/rustc_smir/src/stable_mir/ty.rs
new file mode 100644
index 000000000..f27801b0f
--- /dev/null
+++ b/compiler/rustc_smir/src/stable_mir/ty.rs
@@ -0,0 +1,15 @@
+use super::with;
+
+#[derive(Copy, Clone, Debug)]
+pub struct Ty(pub usize);
+
+impl Ty {
+    pub fn kind(&self) -> TyKind {
+        with(|context| context.ty_kind(*self))
+    }
+}
+
+pub enum TyKind {
+    Bool,
+    Tuple(Vec<Ty>),
+}