std: make TLS accessors closures that return pointers

Author: joboet

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 }
     }
 }