refactor: implement the new SmirContainer struct

Author: makai410

compiler/rustc_smir/src/rustc_internal/mod.rs

@@ -22,10 +22,9 @@ use stable_mir::compiler_interface::SmirInterface;
 use stable_mir::ty::IndexedVal;
 
 use crate::rustc_smir::context::SmirCtxt;
-use crate::rustc_smir::{Stable, Tables};
+use crate::rustc_smir::{Bridge, IndexedVal, SmirContainer, Tables};
 use crate::stable_mir;
 
-mod internal;
 pub mod pretty;
 
 /// Convert an internal Rust compiler item into its stable counterpart, if one exists.
@@ -40,7 +39,7 @@ pub mod pretty;
 ///
 /// This function will panic if StableMIR has not been properly initialized.
 pub fn stable<'tcx, S: Stable<'tcx>>(item: S) -> S::T {
-    with_tables(|tables| item.stable(tables))
+    with_container(|tables, cx| item.stable(tables, cx))
 }
 
 /// Convert a stable item into its internal Rust compiler counterpart, if one exists.
@@ -54,12 +53,11 @@ pub fn stable<'tcx, S: Stable<'tcx>>(item: S) -> S::T {
 /// # Panics
 ///
 /// This function will panic if StableMIR has not been properly initialized.
-pub fn internal<'tcx, S>(tcx: TyCtxt<'tcx>, item: S) -> S::T<'tcx>
+pub fn internal<'tcx, S>(item: S) -> S::T<'tcx>
 where
     S: RustcInternal,
 {
-    // The tcx argument ensures that the item won't outlive the type context.
-    with_tables(|tables| item.internal(tables, tcx))
+    with_container(|tables, cx| item.internal(tables, cx))
 }
 
 impl<'tcx, B: Bridge> Index<B::DefId> for Tables<'tcx, B> {
@@ -101,49 +99,39 @@ pub fn crate_num(item: &stable_mir::Crate) -> CrateNum {
 // datastructures and stable MIR datastructures
 scoped_thread_local! (static TLV: Cell<*const ()>);
 
-pub(crate) fn init<'tcx, F, T>(cx: &SmirCtxt<'tcx>, f: F) -> T
+pub(crate) fn init<'tcx, F, T, B: Bridge>(container: &SmirContainer<'tcx, B>, f: F) -> T
 where
     F: FnOnce() -> T,
 {
     assert!(!TLV.is_set());
-    let ptr = cx as *const _ as *const ();
+    let ptr = container 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<R>(f: impl for<'tcx> FnOnce(&mut Tables<'tcx>) -> R) -> R {
+pub(crate) fn with_container<'tcx, R, B: Bridge>(
+    f: impl FnOnce(&mut Tables<'tcx, B>, &SmirCtxt<'tcx, B>) -> R,
+) -> R {
     assert!(TLV.is_set());
     TLV.with(|tlv| {
         let ptr = tlv.get();
         assert!(!ptr.is_null());
-        let context = ptr as *const SmirCtxt<'_>;
-        let mut tables = unsafe { (*context).0.borrow_mut() };
-        f(&mut *tables)
+        let container = ptr as *const SmirContainer<'tcx, B>;
+        let mut tables = unsafe { (*container).tables.borrow_mut() };
+        let cx = unsafe { (*container).cx.borrow() };
+        f(&mut *tables, &*cx)
     })
 }
 
 pub fn run<F, T>(tcx: TyCtxt<'_>, f: F) -> Result<T, Error>
 where
     F: FnOnce() -> T,
 {
-    let tables = SmirCtxt(RefCell::new(Tables {
-        tcx,
-        def_ids: IndexMap::default(),
-        alloc_ids: IndexMap::default(),
-        spans: IndexMap::default(),
-        types: IndexMap::default(),
-        instances: IndexMap::default(),
-        ty_consts: IndexMap::default(),
-        mir_consts: IndexMap::default(),
-        layouts: IndexMap::default(),
-    }));
-
-    let interface = SmirInterface { cx: tables };
-
-    // Pass the `SmirInterface` to compiler_interface::run
-    // and initialize the rustc-specific TLS with tables.
-    stable_mir::compiler_interface::run(&interface, || init(&interface.cx, f))
+    let smir_cx = RefCell::new(SmirCtxt::new(tcx));
+    let container = SmirContainer { tables: RefCell::new(Tables::new(tcx)), cx: smir_cx };
+
+    stable_mir::compiler_interface::run(&container, || init(&container, f))
 }
 
 /// Instantiate and run the compiler with the provided arguments and callback.

compiler/rustc_smir/src/rustc_smir/mod.rs

@@ -7,21 +7,17 @@
 //!
 //! For now, we are developing everything inside `rustc`, thus, we keep this module private.
 
-use std::ops::RangeInclusive;
+use std::cell::RefCell;
+use std::fmt::Debug;
 
+use context::SmirCtxt;
 use rustc_hir::def::DefKind;
 use rustc_middle::mir;
 use rustc_middle::mir::interpret::AllocId;
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
 use rustc_span::def_id::{CrateNum, DefId, LOCAL_CRATE};
-use stable_mir::abi::Layout;
-use stable_mir::mir::mono::{InstanceDef, StaticDef};
-use stable_mir::ty::{FnDef, MirConstId, Span, TyConstId};
-use stable_mir::{CtorKind, ItemKind};
-use tracing::debug;
 
 use crate::rustc_internal::IndexMap;
-use crate::stable_mir;
 
 mod alloc;
 mod builder;