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(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) }) + }) }