Use scoped-tls for SMIR to map between TyCtxt and SMIR datastructures

4 files changed, 19 insertions, 14 deletions

diff --git a/Cargo.lock b/Cargo.lock
index b78004c93df..1c866b96a95 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4022,6 +4022,7 @@ dependencies = [
  "rustc_hir",
  "rustc_middle",
  "rustc_span",
+ "scoped-tls",
  "tracing",
 ]
 
diff --git a/compiler/rustc_smir/Cargo.toml b/compiler/rustc_smir/Cargo.toml
index 80360a3c73f..a6e6de5f785 100644
--- a/compiler/rustc_smir/Cargo.toml
+++ b/compiler/rustc_smir/Cargo.toml
@@ -8,6 +8,7 @@ rustc_hir = { path = "../rustc_hir" }
 rustc_middle = { path = "../rustc_middle", optional = true }
 rustc_span = { path = "../rustc_span", optional = true }
 tracing = "0.1"
+scoped-tls = "1.0"
 
 [features]
 default = [
diff --git a/compiler/rustc_smir/src/lib.rs b/compiler/rustc_smir/src/lib.rs
index b00f0a1c153..fb03633b99b 100644
--- a/compiler/rustc_smir/src/lib.rs
+++ b/compiler/rustc_smir/src/lib.rs
@@ -19,3 +19,6 @@ pub mod stable_mir;
 
 // Make this module private for now since external users should not call these directly.
 mod rustc_smir;
+
+#[macro_use]
+extern crate scoped_tls;
diff --git a/compiler/rustc_smir/src/stable_mir/mod.rs b/compiler/rustc_smir/src/stable_mir/mod.rs
index 612777b9c75..5e599a77bcd 100644
--- a/compiler/rustc_smir/src/stable_mir/mod.rs
+++ b/compiler/rustc_smir/src/stable_mir/mod.rs
@@ -100,18 +100,17 @@ pub trait Context {
     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()) };
-}
+// 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<*mut ()>);
 
 pub fn run(mut context: impl Context, f: impl FnOnce()) {
-    assert!(TLV.get().is_null());
+    assert!(!TLV.is_set());
     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());
+        let ptr: *mut () = &mut context as *mut &mut _ as _;
+        TLV.set(&Cell::new(ptr), || {
+            f();
+        });
     }
     g(&mut context, f);
 }
@@ -119,9 +118,10 @@ pub fn run(mut context: impl Context, f: impl FnOnce()) {
 /// 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
+    assert!(TLV.is_set());
+    TLV.with(|tlv| {
+        let ptr = tlv.get();
+        assert!(!ptr.is_null());
+        f(unsafe { *(ptr as *mut &mut dyn Context) })
+    })
 }