Auto merge of #123724 - joboet:static_tls, r=m-ou-se

Rewrite TLS on platforms without threads

The saga of #110897 continues!

r? `@m-ou-se` if you have time

2 files changed, 76 insertions, 70 deletions

diff --git a/library/std/src/sys/thread_local/mod.rs b/library/std/src/sys/thread_local/mod.rs
index 6298a5b71cc..36f6f907e72 100644
--- a/library/std/src/sys/thread_local/mod.rs
+++ b/library/std/src/sys/thread_local/mod.rs
@@ -10,7 +10,7 @@ cfg_if::cfg_if! {
         #[doc(hidden)]
         mod static_local;
         #[doc(hidden)]
-        pub use static_local::{Key, thread_local_inner};
+        pub use static_local::{EagerStorage, LazyStorage, thread_local_inner};
     } else if #[cfg(target_thread_local)] {
         #[doc(hidden)]
         mod fast_local;
diff --git a/library/std/src/sys/thread_local/static_local.rs b/library/std/src/sys/thread_local/static_local.rs
index 162c3fbd97a..6beda2e7188 100644
--- a/library/std/src/sys/thread_local/static_local.rs
+++ b/library/std/src/sys/thread_local/static_local.rs
@@ -1,5 +1,7 @@
-use super::lazy::LazyKeyInner;
-use crate::fmt;
+//! On some targets like wasm there's no threads, so no need to generate
+//! thread locals and we can instead just use plain statics!
+
+use crate::cell::UnsafeCell;
 
 #[doc(hidden)]
 #[allow_internal_unstable(thread_local_internals)]
@@ -9,22 +11,17 @@ use crate::fmt;
 pub macro thread_local_inner {
     // used to generate the `LocalKey` value for const-initialized thread locals
     (@key $t:ty, const $init:expr) => {{
-        #[inline] // see comments below
+        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> {
-            const INIT_EXPR: $t = $init;
-
-            // wasm without atomics maps directly to `static mut`, and dtors
-            // aren't implemented because thread dtors aren't really a thing
-            // on wasm right now
-            //
-            // FIXME(#84224) this should come after the `target_thread_local`
-            // block.
-            static mut VAL: $t = INIT_EXPR;
-            // SAFETY: we only ever create shared references, so there's no mutable aliasing.
-            unsafe { $crate::option::Option::Some(&*$crate::ptr::addr_of!(VAL)) }
+            use $crate::thread::local_impl::EagerStorage;
+
+            static VAL: EagerStorage<$t> = EagerStorage { value: __INIT };
+            $crate::option::Option::Some(&VAL.value)
         }
 
         unsafe {
@@ -33,74 +30,83 @@ pub macro thread_local_inner {
     }},
 
     // used to generate the `LocalKey` value for `thread_local!`
-    (@key $t:ty, $init:expr) => {
-        {
-            #[inline]
-            fn __init() -> $t { $init }
-            #[inline]
-            unsafe fn __getit(
-                init: $crate::option::Option<&mut $crate::option::Option<$t>>,
-            ) -> $crate::option::Option<&'static $t> {
-                static __KEY: $crate::thread::local_impl::Key<$t> =
-                    $crate::thread::local_impl::Key::new();
-
-                unsafe {
-                    __KEY.get(move || {
-                        if let $crate::option::Option::Some(init) = init {
-                            if let $crate::option::Option::Some(value) = init.take() {
-                                return value;
-                            } else if $crate::cfg!(debug_assertions) {
-                                $crate::unreachable!("missing default value");
-                            }
-                        }
-                        __init()
-                    })
-                }
-            }
-
-            unsafe {
-                $crate::thread::LocalKey::new(__getit)
-            }
+    (@key $t:ty, $init:expr) => {{
+        #[inline]
+        fn __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::LazyStorage;
+
+            static VAL: LazyStorage<$t> = LazyStorage::new();
+            unsafe { $crate::option::Option::Some(VAL.get(init, __init)) }
         }
-    },
+
+        unsafe {
+            $crate::thread::LocalKey::new(__getit)
+        }
+    }},
     ($(#[$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)*);
     },
 }
 
-/// On some targets like wasm there's no threads, so no need to generate
-/// thread locals and we can instead just use plain statics!
-
-pub struct Key<T> {
-    inner: LazyKeyInner<T>,
+#[allow(missing_debug_implementations)]
+pub struct EagerStorage<T> {
+    pub value: T,
 }
 
-unsafe impl<T> Sync for Key<T> {}
+// SAFETY: the target doesn't have threads.
+unsafe impl<T> Sync for EagerStorage<T> {}
 
-impl<T> fmt::Debug for Key<T> {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_struct("Key").finish_non_exhaustive()
-    }
+#[allow(missing_debug_implementations)]
+pub struct LazyStorage<T> {
+    value: UnsafeCell<Option<T>>,
 }
 
-impl<T> Key<T> {
-    pub const fn new() -> Key<T> {
-        Key { inner: LazyKeyInner::new() }
+impl<T> LazyStorage<T> {
+    pub const fn new() -> LazyStorage<T> {
+        LazyStorage { value: UnsafeCell::new(None) }
     }
 
-    pub unsafe fn get(&self, init: impl FnOnce() -> T) -> Option<&'static T> {
-        // SAFETY: The caller must ensure no reference is ever handed out to
-        // the inner cell nor mutable reference to the Option<T> inside said
-        // cell. This make it safe to hand a reference, though the lifetime
-        // of 'static is itself unsafe, making the get method unsafe.
-        let value = unsafe {
-            match self.inner.get() {
-                Some(ref value) => value,
-                None => self.inner.initialize(init),
-            }
-        };
-
-        Some(value)
+    /// Gets a reference to the contained value, initializing it if necessary.
+    ///
+    /// # Safety
+    /// The returned reference may not be used after reentrant initialization has occurred.
+    #[inline]
+    pub unsafe fn get(
+        &'static self,
+        i: Option<&mut Option<T>>,
+        f: impl FnOnce() -> T,
+    ) -> &'static T {
+        let value = unsafe { &*self.value.get() };
+        match value {
+            Some(v) => v,
+            None => self.initialize(i, f),
+        }
+    }
+
+    #[cold]
+    unsafe fn initialize(
+        &'static self,
+        i: Option<&mut Option<T>>,
+        f: impl FnOnce() -> T,
+    ) -> &'static T {
+        let value = i.and_then(Option::take).unwrap_or_else(f);
+        // Destroy the old value, after updating the TLS variable as the
+        // destructor might reference it.
+        // FIXME(#110897): maybe panic on recursive initialization.
+        unsafe {
+            self.value.get().replace(Some(value));
+        }
+        // SAFETY: we just set this to `Some`.
+        unsafe { (*self.value.get()).as_ref().unwrap_unchecked() }
     }
 }
+
+// SAFETY: the target doesn't have threads.
+unsafe impl<T> Sync for LazyStorage<T> {}