2 files changed, 251 insertions, 91 deletions

diff --git a/compiler/rustc_smir/src/rustc_internal/internal.rs b/compiler/rustc_smir/src/rustc_internal/internal.rs
new file mode 100644
index 000000000..7cfdbbbf7
--- /dev/null
+++ b/compiler/rustc_smir/src/rustc_internal/internal.rs
@@ -0,0 +1,67 @@
+//! Module containing the translation from stable mir constructs to the rustc counterpart.
+//!
+//! This module will only include a few constructs to allow users to invoke internal rustc APIs
+//! due to incomplete stable coverage.
+
+// Prefer importing stable_mir over internal rustc constructs to make this file more readable.
+use crate::rustc_smir::Tables;
+use rustc_middle::ty::{self as rustc_ty, Ty as InternalTy};
+use stable_mir::ty::{Const, GenericArgKind, GenericArgs, Region, Ty};
+use stable_mir::DefId;
+
+use super::RustcInternal;
+
+impl<'tcx> RustcInternal<'tcx> for DefId {
+    type T = rustc_span::def_id::DefId;
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        tables.def_ids[*self]
+    }
+}
+
+impl<'tcx> RustcInternal<'tcx> for GenericArgs {
+    type T = rustc_ty::GenericArgsRef<'tcx>;
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        tables.tcx.mk_args_from_iter(self.0.iter().map(|arg| arg.internal(tables)))
+    }
+}
+
+impl<'tcx> RustcInternal<'tcx> for GenericArgKind {
+    type T = rustc_ty::GenericArg<'tcx>;
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        match self {
+            GenericArgKind::Lifetime(reg) => reg.internal(tables).into(),
+            GenericArgKind::Type(ty) => ty.internal(tables).into(),
+            GenericArgKind::Const(cnst) => ty_const(cnst, tables).into(),
+        }
+    }
+}
+
+impl<'tcx> RustcInternal<'tcx> for Region {
+    type T = rustc_ty::Region<'tcx>;
+    fn internal(&self, _tables: &mut Tables<'tcx>) -> Self::T {
+        todo!()
+    }
+}
+
+impl<'tcx> RustcInternal<'tcx> for Ty {
+    type T = InternalTy<'tcx>;
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        tables.types[*self]
+    }
+}
+
+fn ty_const<'tcx>(constant: &Const, tables: &mut Tables<'tcx>) -> rustc_ty::Const<'tcx> {
+    match constant.internal(tables) {
+        rustc_middle::mir::Const::Ty(c) => c,
+        cnst => {
+            panic!("Trying to convert constant `{constant:?}` to type constant, but found {cnst:?}")
+        }
+    }
+}
+
+impl<'tcx> RustcInternal<'tcx> for Const {
+    type T = rustc_middle::mir::Const<'tcx>;
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        tables.constants[self.id]
+    }
+}
diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index 1a9dea99f..4d2a51822 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -3,24 +3,38 @@
 //! For that, we define APIs that will temporarily be public to 3P that exposes rustc internal APIs
 //! until stable MIR is complete.
 
-use std::ops::{ControlFlow, Index};
-
-use crate::rustc_internal;
-use crate::rustc_smir::Tables;
-use rustc_driver::{Callbacks, Compilation, RunCompiler};
-use rustc_interface::{interface, Queries};
+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;
-pub use rustc_span::def_id::{CrateNum, DefId};
+use rustc_span::def_id::{CrateNum, DefId};
 use rustc_span::Span;
-use stable_mir::CompilerError;
+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;
 
     #[inline(always)]
     fn index(&self, index: stable_mir::DefId) -> &Self::Output {
-        &self.def_ids[index.0]
+        &self.def_ids[index]
     }
 }
 
@@ -29,7 +43,7 @@ impl<'tcx> Index<stable_mir::ty::Span> for Tables<'tcx> {
 
     #[inline(always)]
     fn index(&self, index: stable_mir::ty::Span) -> &Self::Output {
-        &self.spans[index.0]
+        &self.spans[index]
     }
 }
 
@@ -54,8 +68,8 @@ impl<'tcx> Tables<'tcx> {
         stable_mir::ty::ClosureDef(self.create_def_id(did))
     }
 
-    pub fn generator_def(&mut self, did: DefId) -> stable_mir::ty::GeneratorDef {
-        stable_mir::ty::GeneratorDef(self.create_def_id(did))
+    pub fn coroutine_def(&mut self, did: DefId) -> stable_mir::ty::CoroutineDef {
+        stable_mir::ty::CoroutineDef(self.create_def_id(did))
     }
 
     pub fn alias_def(&mut self, did: DefId) -> stable_mir::ty::AliasDef {
@@ -94,37 +108,27 @@ impl<'tcx> Tables<'tcx> {
         stable_mir::ty::Prov(self.create_alloc_id(aid))
     }
 
-    fn create_def_id(&mut self, did: DefId) -> stable_mir::DefId {
-        // 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::DefId(i);
-            }
-        }
-        let id = self.def_ids.len();
-        self.def_ids.push(did);
-        stable_mir::DefId(id)
+    pub(crate) fn create_def_id(&mut self, did: DefId) -> stable_mir::DefId {
+        self.def_ids.create_or_fetch(did)
     }
 
     fn create_alloc_id(&mut self, aid: AllocId) -> stable_mir::AllocId {
-        // FIXME: this becomes inefficient when we have too many ids
-        if let Some(i) = self.alloc_ids.iter().position(|a| *a == aid) {
-            return stable_mir::AllocId(i);
-        };
-        let id = self.def_ids.len();
-        self.alloc_ids.push(aid);
-        stable_mir::AllocId(id)
+        self.alloc_ids.create_or_fetch(aid)
     }
 
     pub(crate) fn create_span(&mut self, span: Span) -> stable_mir::ty::Span {
-        for (i, &sp) in self.spans.iter().enumerate() {
-            if sp == span {
-                return stable_mir::ty::Span(i);
-            }
-        }
-        let id = self.spans.len();
-        self.spans.push(span);
-        stable_mir::ty::Span(id)
+        self.spans.create_or_fetch(span)
+    }
+
+    pub(crate) fn instance_def(
+        &mut self,
+        instance: ty::Instance<'tcx>,
+    ) -> stable_mir::mir::mono::InstanceDef {
+        self.instances.create_or_fetch(instance)
+    }
+
+    pub(crate) fn static_def(&mut self, did: DefId) -> stable_mir::mir::mono::StaticDef {
+        stable_mir::mir::mono::StaticDef(self.create_def_id(did))
     }
 }
 
@@ -132,68 +136,157 @@ pub fn crate_num(item: &stable_mir::Crate) -> CrateNum {
     item.id.into()
 }
 
-pub fn run(tcx: TyCtxt<'_>, f: impl FnOnce()) {
-    stable_mir::run(
-        Tables { tcx, def_ids: vec![], alloc_ids: vec![], spans: vec![], types: vec![] },
-        f,
-    );
+// 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();
+    });
 }
 
-pub struct StableMir<B = (), C = ()>
-where
-    B: Send,
-    C: Send,
-{
-    args: Vec<String>,
-    callback: fn(TyCtxt<'_>) -> ControlFlow<B, C>,
-    result: Option<ControlFlow<B, C>>,
+/// 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)
+    })
 }
 
-impl<B, C> StableMir<B, C>
-where
-    B: Send,
-    C: Send,
-{
-    /// Creates a new `StableMir` instance, with given test_function and arguments.
-    pub fn new(args: Vec<String>, callback: fn(TyCtxt<'_>) -> ControlFlow<B, C>) -> Self {
-        StableMir { args, callback, result: None }
-    }
-
-    /// Runs the compiler against given target and tests it with `test_function`
-    pub fn run(&mut self) -> Result<C, CompilerError<B>> {
-        let compiler_result =
-            rustc_driver::catch_fatal_errors(|| RunCompiler::new(&self.args.clone(), self).run());
-        match (compiler_result, self.result.take()) {
-            (Ok(Ok(())), Some(ControlFlow::Continue(value))) => Ok(value),
-            (Ok(Ok(())), Some(ControlFlow::Break(value))) => Err(CompilerError::Interrupted(value)),
-            (Ok(Ok(_)), None) => Err(CompilerError::Skipped),
-            (Ok(Err(_)), _) => Err(CompilerError::CompilationFailed),
-            (Err(_), _) => Err(CompilerError::ICE),
+pub fn run(tcx: TyCtxt<'_>, f: impl FnOnce()) {
+    let tables = TablesWrapper(RefCell::new(Tables {
+        tcx,
+        def_ids: IndexMap::default(),
+        alloc_ids: IndexMap::default(),
+        spans: IndexMap::default(),
+        types: IndexMap::default(),
+        instances: IndexMap::default(),
+        constants: IndexMap::default(),
+    }));
+    stable_mir::run(&tables, || init(&tables, f));
+}
+
+#[macro_export]
+macro_rules! run {
+    ($args:expr, $callback:expr) => {
+        run!($args, tcx, $callback)
+    };
+    ($args:expr, $tcx:ident, $callback:expr) => {{
+        use rustc_driver::{Callbacks, Compilation, RunCompiler};
+        use rustc_interface::{interface, Queries};
+        use stable_mir::CompilerError;
+        use std::ops::ControlFlow;
+
+        pub struct StableMir<B = (), C = ()>
+        where
+            B: Send,
+            C: Send,
+        {
+            args: Vec<String>,
+            callback: fn(TyCtxt<'_>) -> ControlFlow<B, C>,
+            result: Option<ControlFlow<B, C>>,
+        }
+
+        impl<B, C> StableMir<B, C>
+        where
+            B: Send,
+            C: Send,
+        {
+            /// Creates a new `StableMir` instance, with given test_function and arguments.
+            pub fn new(args: Vec<String>, callback: fn(TyCtxt<'_>) -> ControlFlow<B, C>) -> Self {
+                StableMir { args, callback, result: None }
+            }
+
+            /// Runs the compiler against given target and tests it with `test_function`
+            pub fn run(&mut self) -> Result<C, CompilerError<B>> {
+                let compiler_result = rustc_driver::catch_fatal_errors(|| {
+                    RunCompiler::new(&self.args.clone(), self).run()
+                });
+                match (compiler_result, self.result.take()) {
+                    (Ok(Ok(())), Some(ControlFlow::Continue(value))) => Ok(value),
+                    (Ok(Ok(())), Some(ControlFlow::Break(value))) => {
+                        Err(CompilerError::Interrupted(value))
+                    }
+                    (Ok(Ok(_)), None) => Err(CompilerError::Skipped),
+                    (Ok(Err(_)), _) => Err(CompilerError::CompilationFailed),
+                    (Err(_), _) => Err(CompilerError::ICE),
+                }
+            }
         }
+
+        impl<B, C> Callbacks for StableMir<B, C>
+        where
+            B: Send,
+            C: Send,
+        {
+            /// Called after analysis. Return value instructs the compiler whether to
+            /// continue the compilation afterwards (defaults to `Compilation::Continue`)
+            fn after_analysis<'tcx>(
+                &mut self,
+                _compiler: &interface::Compiler,
+                queries: &'tcx Queries<'tcx>,
+            ) -> Compilation {
+                queries.global_ctxt().unwrap().enter(|tcx| {
+                    rustc_internal::run(tcx, || {
+                        self.result = Some((self.callback)(tcx));
+                    });
+                    if self.result.as_ref().is_some_and(|val| val.is_continue()) {
+                        Compilation::Continue
+                    } else {
+                        Compilation::Stop
+                    }
+                })
+            }
+        }
+
+        StableMir::new($args, |$tcx| $callback).run()
+    }};
+}
+
+/// Simmilar to rustc's `FxIndexMap`, `IndexMap` with extra
+/// safety features added.
+pub struct IndexMap<K, V> {
+    index_map: fx::FxIndexMap<K, V>,
+}
+
+impl<K, V> Default for IndexMap<K, V> {
+    fn default() -> Self {
+        Self { index_map: FxIndexMap::default() }
+    }
+}
+
+impl<K: PartialEq + Hash + Eq, V: Copy + Debug + PartialEq + IndexedVal> IndexMap<K, V> {
+    pub fn create_or_fetch(&mut self, key: K) -> V {
+        let len = self.index_map.len();
+        let v = self.index_map.entry(key).or_insert(V::to_val(len));
+        *v
     }
 }
 
-impl<B, C> Callbacks for StableMir<B, C>
-where
-    B: Send,
-    C: Send,
+impl<K: PartialEq + Hash + Eq, V: Copy + Debug + PartialEq + IndexedVal> Index<V>
+    for IndexMap<K, V>
 {
-    /// Called after analysis. Return value instructs the compiler whether to
-    /// continue the compilation afterwards (defaults to `Compilation::Continue`)
-    fn after_analysis<'tcx>(
-        &mut self,
-        _compiler: &interface::Compiler,
-        queries: &'tcx Queries<'tcx>,
-    ) -> Compilation {
-        queries.global_ctxt().unwrap().enter(|tcx| {
-            rustc_internal::run(tcx, || {
-                self.result = Some((self.callback)(tcx));
-            });
-            if self.result.as_ref().is_some_and(|val| val.is_continue()) {
-                Compilation::Continue
-            } else {
-                Compilation::Stop
-            }
-        })
+    type Output = K;
+
+    fn index(&self, index: V) -> &Self::Output {
+        let (k, v) = self.index_map.get_index(index.to_index()).unwrap();
+        assert_eq!(*v, index, "Provided value doesn't match with indexed value");
+        k
     }
 }
+
+/// Trait used to translate a stable construct to its rustc counterpart.
+///
+/// This is basically a mirror of [crate::rustc_smir::Stable].
+pub trait RustcInternal<'tcx> {
+    type T;
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T;
+}