Auto merge of rust-lang#125525 - joboet:tls_accessor, r=cuviper

Make TLS accessors closures that return pointers

The current TLS macros generate a function that returns an `Option<&'static T>`. This is both risky as we lie about lifetimes, and necessitates that those functions are `unsafe`. By returning a `*const T` instead, the accessor function do not have safety requirements any longer and can be made closures without hassle. This PR does exactly that!

For native TLS, the closure approach makes it trivial to select the right accessor function at compile-time, which could result in a slight speed-up (I have the hope that the accessors are now simple enough for the MIR-inliner to kick in).

Author: bors

std/src/sys/thread_local/fast_local/eager.rs

@@ -21,43 +21,35 @@ impl<T> Storage<T> {
         Storage { state: Cell::new(State::Initial), val: UnsafeCell::new(val) }
     }
 
-    /// Get a reference to the TLS value. If the TLS variable has been destroyed,
-    /// `None` is returned.
+    /// Get a pointer to the TLS value. If the TLS variable has been destroyed,
+    /// a null pointer is returned.
     ///
-    /// # Safety
-    /// * The `self` reference must remain valid until the TLS destructor has been
-    ///   run.
-    /// * The returned reference may only be used until thread destruction occurs
-    ///   and may not be used after reentrant initialization has occurred.
+    /// The resulting pointer may not be used after thread destruction has
+    /// occurred.
     ///
-    // FIXME(#110897): return NonNull instead of lying about the lifetime.
+    /// # Safety
+    /// The `self` reference must remain valid until the TLS destructor is run.
     #[inline]
-    pub unsafe fn get(&self) -> Option<&'static T> {
+    pub unsafe fn get(&self) -> *const T {
         match self.state.get() {
-            // SAFETY: as the state is not `Destroyed`, the value cannot have
-            // been destroyed yet. The reference fulfills the terms outlined
-            // above.
-            State::Alive => unsafe { Some(&*self.val.get()) },
-            State::Destroyed => None,
+            State::Alive => self.val.get(),
+            State::Destroyed => ptr::null(),
             State::Initial => unsafe { self.initialize() },
         }
     }
 
     #[cold]
-    unsafe fn initialize(&self) -> Option<&'static T> {
+    unsafe fn initialize(&self) -> *const T {
         // Register the destructor
 
         // SAFETY:
-        // * the destructor will be called at thread destruction.
-        // * the caller guarantees that `self` will be valid until that time.
+        // The caller guarantees that `self` will be valid until thread destruction.
         unsafe {
             register_dtor(ptr::from_ref(self).cast_mut().cast(), destroy::<T>);
         }
+
         self.state.set(State::Alive);
-        // SAFETY: as the state is not `Destroyed`, the value cannot have
-        // been destroyed yet. The reference fulfills the terms outlined
-        // above.
-        unsafe { Some(&*self.val.get()) }
+        self.val.get()
     }
 }
 

std/src/sys/thread_local/fast_local/lazy.rs

@@ -39,49 +39,31 @@ where
         Storage { state: UnsafeCell::new(State::Initial) }
     }
 
-    /// Get a reference to the TLS value, potentially initializing it with the
-    /// provided parameters. If the TLS variable has been destroyed, `None` is
-    /// returned.
+    /// Get a pointer to the TLS value, potentially initializing it with the
+    /// provided parameters. If the TLS variable has been destroyed, a null
+    /// pointer is returned.
     ///
-    /// # Safety
-    /// * The `self` reference must remain valid until the TLS destructor is run,
-    ///   at which point the returned reference is invalidated.
-    /// * The returned reference may only be used until thread destruction occurs
-    ///   and may not be used after reentrant initialization has occurred.
+    /// The resulting pointer may not be used after reentrant inialialization
+    /// or thread destruction has occurred.
     ///
-    // FIXME(#110897): return NonNull instead of lying about the lifetime.
+    /// # Safety
+    /// The `self` reference must remain valid until the TLS destructor is run.
     #[inline]
-    pub unsafe fn get_or_init(
-        &self,
-        i: Option<&mut Option<T>>,
-        f: impl FnOnce() -> T,
-    ) -> Option<&'static T> {
-        // SAFETY:
-        // No mutable reference to the inner value exists outside the calls to
-        // `replace`. The lifetime of the returned reference fulfills the terms
-        // outlined above.
+    pub unsafe fn get_or_init(&self, i: Option<&mut Option<T>>, f: impl FnOnce() -> T) -> *const T {
         let state = unsafe { &*self.state.get() };
         match state {
-            State::Alive(v) => Some(v),
-            State::Destroyed(_) => None,
+            State::Alive(v) => v,
+            State::Destroyed(_) => ptr::null(),
             State::Initial => unsafe { self.initialize(i, f) },
         }
     }
 
     #[cold]
-    unsafe fn initialize(
-        &self,
-        i: Option<&mut Option<T>>,
-        f: impl FnOnce() -> T,
-    ) -> Option<&'static T> {
+    unsafe fn initialize(&self, i: Option<&mut Option<T>>, f: impl FnOnce() -> T) -> *const T {
         // Perform initialization
 
         let v = i.and_then(Option::take).unwrap_or_else(f);
 
-        // SAFETY:
-        // If references to the inner value exist, they were created in `f`
-        // and are invalidated here. The caller promises to never use them
-        // after this.
         let old = unsafe { self.state.get().replace(State::Alive(v)) };
         match old {
             // If the variable is not being recursively initialized, register
@@ -92,12 +74,10 @@ where
             val => drop(val),
         }
 
-        // SAFETY:
-        // Initialization was completed and the state was set to `Alive`, so the
-        // reference fulfills the terms outlined above.
+        // SAFETY: the state was just set to `Alive`
         unsafe {
             let State::Alive(v) = &*self.state.get() else { unreachable_unchecked() };
-            Some(v)
+            v
         }
     }
 }

std/src/sys/thread_local/fast_local/mod.rs

@@ -52,32 +52,26 @@ pub macro thread_local_inner {
     (@key $t:ty, const $init:expr) => {{
         const __INIT: $t = $init;
 
-        #[inline]
-        #[deny(unsafe_op_in_unsafe_fn)]
-        unsafe fn __getit(
-            _init: $crate::option::Option<&mut $crate::option::Option<$t>>,
-        ) -> $crate::option::Option<&'static $t> {
-            use $crate::thread::local_impl::EagerStorage;
+        unsafe {
             use $crate::mem::needs_drop;
-            use $crate::ptr::addr_of;
+            use $crate::thread::LocalKey;
+            use $crate::thread::local_impl::EagerStorage;
 
-            if needs_drop::<$t>() {
-                #[thread_local]
-                static VAL: EagerStorage<$t> = EagerStorage::new(__INIT);
-                unsafe {
-                    VAL.get()
+            LocalKey::new(const {
+                if needs_drop::<$t>() {
+                    |_| {
+                        #[thread_local]
+                        static VAL: EagerStorage<$t> = EagerStorage::new(__INIT);
+                        VAL.get()
+                    }
+                } else {
+                    |_| {
+                        #[thread_local]
+                        static VAL: $t = __INIT;
+                        &VAL
+                    }
                 }
-            } else {
-                #[thread_local]
-                static VAL: $t = __INIT;
-                unsafe {
-                    $crate::option::Option::Some(&*addr_of!(VAL))
-                }
-            }
-        }
-
-        unsafe {
-            $crate::thread::LocalKey::new(__getit)
+            })
         }
     }},
 
@@ -88,31 +82,26 @@ pub macro thread_local_inner {
             $init
         }
 
-        #[inline]
-        #[deny(unsafe_op_in_unsafe_fn)]
-        unsafe fn __getit(
-            init: $crate::option::Option<&mut $crate::option::Option<$t>>,
-        ) -> $crate::option::Option<&'static $t> {
-            use $crate::thread::local_impl::LazyStorage;
+        unsafe {
             use $crate::mem::needs_drop;
+            use $crate::thread::LocalKey;
+            use $crate::thread::local_impl::LazyStorage;
 
-            if needs_drop::<$t>() {
-                #[thread_local]
-                static VAL: LazyStorage<$t, ()> = LazyStorage::new();
-                unsafe {
-                    VAL.get_or_init(init, __init)
+            LocalKey::new(const {
+                if needs_drop::<$t>() {
+                    |init| {
+                        #[thread_local]
+                        static VAL: LazyStorage<$t, ()> = LazyStorage::new();
+                        VAL.get_or_init(init, __init)
+                    }
+                } else {
+                    |init| {
+                        #[thread_local]
+                        static VAL: LazyStorage<$t, !> = LazyStorage::new();
+                        VAL.get_or_init(init, __init)
+                    }
                 }
-            } else {
-                #[thread_local]
-                static VAL: LazyStorage<$t, !> = LazyStorage::new();
-                unsafe {
-                    VAL.get_or_init(init, __init)
-                }
-            }
-        }
-
-        unsafe {
-            $crate::thread::LocalKey::new(__getit)
+            })
         }
     }},
     ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty, $($init:tt)*) => {

std/src/sys/thread_local/os_local.rs

@@ -16,30 +16,22 @@ pub macro thread_local_inner {
     },
 
     // used to generate the `LocalKey` value for `thread_local!`
-    (@key $t:ty, $init:expr) => {
-        {
-            #[inline]
-            fn __init() -> $t { $init }
+    (@key $t:ty, $init:expr) => {{
+        #[inline]
+        fn __init() -> $t { $init }
 
-            // `#[inline] does not work on windows-gnu due to linking errors around dllimports.
-            // See https://github.com/rust-lang/rust/issues/109797.
-            #[cfg_attr(not(windows), inline)]
-            unsafe fn __getit(
-                init: $crate::option::Option<&mut $crate::option::Option<$t>>,
-            ) -> $crate::option::Option<&'static $t> {
-                use $crate::thread::local_impl::Key;
-
-                static __KEY: Key<$t> = Key::new();
-                unsafe {
-                    __KEY.get(init, __init)
-                }
-            }
+        unsafe {
+            use $crate::thread::LocalKey;
+            use $crate::thread::local_impl::Key;
 
-            unsafe {
-                $crate::thread::LocalKey::new(__getit)
-            }
+            // Inlining does not work on windows-gnu due to linking errors around
+            // dllimports. See https://github.com/rust-lang/rust/issues/109797.
+            LocalKey::new(#[cfg_attr(windows, inline(never))] |init| {
+                static VAL: Key<$t> = Key::new();
+                VAL.get(init, __init)
+            })
         }
-    },
+    }},
     ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty, $($init:tt)*) => {
         $(#[$attr])* $vis const $name: $crate::thread::LocalKey<$t> =
             $crate::thread::local_impl::thread_local_inner!(@key $t, $($init)*);
@@ -67,38 +59,33 @@ impl<T: 'static> Key<T> {
         Key { os: OsKey::new(Some(destroy_value::<T>)), marker: PhantomData }
     }
 
-    /// Get the value associated with this key, initializating it if necessary.
+    /// Get a pointer to the TLS value, potentially initializing it with the
+    /// provided parameters. If the TLS variable has been destroyed, a null
+    /// pointer is returned.
     ///
-    /// # Safety
-    /// * the returned reference must not be used after recursive initialization
-    /// or thread destruction occurs.
-    pub unsafe fn get(
-        &'static self,
-        i: Option<&mut Option<T>>,
-        f: impl FnOnce() -> T,
-    ) -> Option<&'static T> {
+    /// The resulting pointer may not be used after reentrant inialialization
+    /// or thread destruction has occurred.
+    pub fn get(&'static self, i: Option<&mut Option<T>>, f: impl FnOnce() -> T) -> *const T {
         // SAFETY: (FIXME: get should actually be safe)
         let ptr = unsafe { self.os.get() as *mut Value<T> };
         if ptr.addr() > 1 {
             // SAFETY: the check ensured the pointer is safe (its destructor
             // is not running) + it is coming from a trusted source (self).
-            unsafe { Some(&(*ptr).value) }
+            unsafe { &(*ptr).value }
         } else {
-            // SAFETY: At this point we are sure we have no value and so
-            // initializing (or trying to) is safe.
-            unsafe { self.try_initialize(ptr, i, f) }
+            self.try_initialize(ptr, i, f)
         }
     }
 
-    unsafe fn try_initialize(
+    fn try_initialize(
         &'static self,
         ptr: *mut Value<T>,
         i: Option<&mut Option<T>>,
         f: impl FnOnce() -> T,
-    ) -> Option<&'static T> {
+    ) -> *const T {
         if ptr.addr() == 1 {
             // destructor is running
-            return None;
+            return ptr::null();
         }
 
         let value = i.and_then(Option::take).unwrap_or_else(f);
@@ -119,7 +106,7 @@ impl<T: 'static> Key<T> {
         }
 
         // SAFETY: We just created this value above.
-        unsafe { Some(&(*ptr).value) }
+        unsafe { &(*ptr).value }
     }
 }