std: refactor the TLS implementation

As discovered by Mara in #110897, our TLS implementation is a total mess. In the past months, I have simplified the actual macros and their expansions, but the majority of the complexity comes from the platform-specific support code needed to create keys and register destructors. In keeping with #117276, I have therefore moved all of the `thread_local_key`/`thread_local_dtor` modules to the `thread_local` module in `sys` and merged them into a new structure, so that future porters of `std` can simply mix-and-match the existing code instead of having to copy the same (bad) implementation everywhere. The new structure should become obvious when looking at `sys/thread_local/mod.rs`.

Unfortunately, the documentation changes associated with the refactoring have made this PR rather large. That said, this contains no functional changes except for two small ones:
* the key-based destructor fallback now, by virtue of sharing the implementation used by macOS and others, stores its list in a `#[thread_local]` static instead of in the key, eliminating one indirection layer and drastically simplifying its code.
* I've switched over ZKVM (tier 3) to use the same implementation as WebAssembly, as the implementation was just a way worse version of that

Please let me know if I can make this easier to review! I know these large PRs aren't optimal, but I couldn't think of any good intermediate steps.

@rustbot label +A-thread-locals

3 files changed, 284 insertions, 0 deletions

diff --git a/library/std/src/sys/thread_local/native/eager.rs b/library/std/src/sys/thread_local/native/eager.rs
new file mode 100644
index 00000000000..99e5ae7fb96
--- /dev/null
+++ b/library/std/src/sys/thread_local/native/eager.rs
@@ -0,0 +1,74 @@
+use crate::cell::{Cell, UnsafeCell};
+use crate::ptr::{self, drop_in_place};
+use crate::sys::thread_local::abort_on_dtor_unwind;
+use crate::sys::thread_local::destructors;
+
+#[derive(Clone, Copy)]
+enum State {
+    Initial,
+    Alive,
+    Destroyed,
+}
+
+#[allow(missing_debug_implementations)]
+pub struct Storage<T> {
+    state: Cell<State>,
+    val: UnsafeCell<T>,
+}
+
+impl<T> Storage<T> {
+    pub const fn new(val: T) -> Storage<T> {
+        Storage { state: Cell::new(State::Initial), val: UnsafeCell::new(val) }
+    }
+
+    /// Get a pointer to the TLS value. If the TLS variable has been destroyed,
+    /// a null pointer is returned.
+    ///
+    /// The resulting pointer may not be used after thread destruction has
+    /// occurred.
+    ///
+    /// # Safety
+    /// The `self` reference must remain valid until the TLS destructor is run.
+    #[inline]
+    pub unsafe fn get(&self) -> *const T {
+        match self.state.get() {
+            State::Alive => self.val.get(),
+            State::Destroyed => ptr::null(),
+            State::Initial => unsafe { self.initialize() },
+        }
+    }
+
+    #[cold]
+    unsafe fn initialize(&self) -> *const T {
+        // Register the destructor
+
+        // SAFETY:
+        // The caller guarantees that `self` will be valid until thread destruction.
+        unsafe {
+            destructors::register(ptr::from_ref(self).cast_mut().cast(), destroy::<T>);
+        }
+
+        self.state.set(State::Alive);
+        self.val.get()
+    }
+}
+
+/// Transition an `Alive` TLS variable into the `Destroyed` state, dropping its
+/// value.
+///
+/// # Safety
+/// * Must only be called at thread destruction.
+/// * `ptr` must point to an instance of `Storage` with `Alive` state and be
+///   valid for accessing that instance.
+unsafe extern "C" fn destroy<T>(ptr: *mut u8) {
+    // Print a nice abort message if a panic occurs.
+    abort_on_dtor_unwind(|| {
+        let storage = unsafe { &*(ptr as *const Storage<T>) };
+        // Update the state before running the destructor as it may attempt to
+        // access the variable.
+        storage.state.set(State::Destroyed);
+        unsafe {
+            drop_in_place(storage.val.get());
+        }
+    })
+}
diff --git a/library/std/src/sys/thread_local/native/lazy.rs b/library/std/src/sys/thread_local/native/lazy.rs
new file mode 100644
index 00000000000..9d47e8ef689
--- /dev/null
+++ b/library/std/src/sys/thread_local/native/lazy.rs
@@ -0,0 +1,101 @@
+use crate::cell::UnsafeCell;
+use crate::hint::unreachable_unchecked;
+use crate::ptr;
+use crate::sys::thread_local::abort_on_dtor_unwind;
+use crate::sys::thread_local::destructors;
+
+pub unsafe trait DestroyedState: Sized {
+    fn register_dtor<T>(s: &Storage<T, Self>);
+}
+
+unsafe impl DestroyedState for ! {
+    fn register_dtor<T>(_: &Storage<T, !>) {}
+}
+
+unsafe impl DestroyedState for () {
+    fn register_dtor<T>(s: &Storage<T, ()>) {
+        unsafe {
+            destructors::register(ptr::from_ref(s).cast_mut().cast(), destroy::<T>);
+        }
+    }
+}
+
+enum State<T, D> {
+    Initial,
+    Alive(T),
+    Destroyed(D),
+}
+
+#[allow(missing_debug_implementations)]
+pub struct Storage<T, D> {
+    state: UnsafeCell<State<T, D>>,
+}
+
+impl<T, D> Storage<T, D>
+where
+    D: DestroyedState,
+{
+    pub const fn new() -> Storage<T, D> {
+        Storage { state: UnsafeCell::new(State::Initial) }
+    }
+
+    /// 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.
+    ///
+    /// The resulting pointer may not be used after reentrant inialialization
+    /// or thread destruction has occurred.
+    ///
+    /// # 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) -> *const T {
+        let state = unsafe { &*self.state.get() };
+        match state {
+            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) -> *const T {
+        // Perform initialization
+
+        let v = i.and_then(Option::take).unwrap_or_else(f);
+
+        let old = unsafe { self.state.get().replace(State::Alive(v)) };
+        match old {
+            // If the variable is not being recursively initialized, register
+            // the destructor. This might be a noop if the value does not need
+            // destruction.
+            State::Initial => D::register_dtor(self),
+            // Else, drop the old value. This might be changed to a panic.
+            val => drop(val),
+        }
+
+        // SAFETY: the state was just set to `Alive`
+        unsafe {
+            let State::Alive(v) = &*self.state.get() else { unreachable_unchecked() };
+            v
+        }
+    }
+}
+
+/// Transition an `Alive` TLS variable into the `Destroyed` state, dropping its
+/// value.
+///
+/// # Safety
+/// * Must only be called at thread destruction.
+/// * `ptr` must point to an instance of `Storage<T, ()>` and be valid for
+///   accessing that instance.
+unsafe extern "C" fn destroy<T>(ptr: *mut u8) {
+    // Print a nice abort message if a panic occurs.
+    abort_on_dtor_unwind(|| {
+        let storage = unsafe { &*(ptr as *const Storage<T, ()>) };
+        // Update the state before running the destructor as it may attempt to
+        // access the variable.
+        let val = unsafe { storage.state.get().replace(State::Destroyed(())) };
+        drop(val);
+    })
+}
diff --git a/library/std/src/sys/thread_local/native/mod.rs b/library/std/src/sys/thread_local/native/mod.rs
new file mode 100644
index 00000000000..1cc45fe892d
--- /dev/null
+++ b/library/std/src/sys/thread_local/native/mod.rs
@@ -0,0 +1,109 @@
+//! Thread local support for platforms with native TLS.
+//!
+//! To achieve the best performance, we choose from four different types for
+//! the TLS variable, depending on the method of initialization used (`const`
+//! or lazy) and the drop requirements of the stored type:
+//!
+//! |         | `Drop`               | `!Drop`             |
+//! |--------:|:--------------------:|:-------------------:|
+//! | `const` | `EagerStorage<T>`    | `T`                 |
+//! | lazy    | `LazyStorage<T, ()>` | `LazyStorage<T, !>` |
+//!
+//! For `const` initialization and `!Drop` types, we simply use `T` directly,
+//! but for other situations, we implement a state machine to handle
+//! initialization of the variable and its destructor and destruction.
+//! Upon accessing the TLS variable, the current state is compared:
+//!
+//! 1. If the state is `Initial`, initialize the storage, transition the state
+//!    to `Alive` and (if applicable) register the destructor, and return a
+//!    reference to the value.
+//! 2. If the state is `Alive`, initialization was previously completed, so
+//!    return a reference to the value.
+//! 3. If the state is `Destroyed`, the destructor has been run already, so
+//!    return [`None`].
+//!
+//! The TLS destructor sets the state to `Destroyed` and drops the current value.
+//!
+//! To simplify the code, we make `LazyStorage` generic over the destroyed state
+//! and use the `!` type (never type) as type parameter for `!Drop` types. This
+//! eliminates the `Destroyed` state for these values, which can allow more niche
+//! optimizations to occur for the `State` enum. For `Drop` types, `()` is used.
+
+mod eager;
+mod lazy;
+
+pub use eager::Storage as EagerStorage;
+pub use lazy::Storage as LazyStorage;
+
+#[doc(hidden)]
+#[allow_internal_unstable(
+    thread_local_internals,
+    cfg_target_thread_local,
+    thread_local,
+    never_type
+)]
+#[allow_internal_unsafe]
+#[unstable(feature = "thread_local_internals", issue = "none")]
+#[rustc_macro_transparency = "semitransparent"]
+pub macro thread_local_inner {
+    // used to generate the `LocalKey` value for const-initialized thread locals
+    (@key $t:ty, const $init:expr) => {{
+        const __INIT: $t = $init;
+
+        unsafe {
+            use $crate::mem::needs_drop;
+            use $crate::thread::LocalKey;
+            use $crate::thread::local_impl::EagerStorage;
+
+            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
+                    }
+                }
+            })
+        }
+    }},
+
+    // used to generate the `LocalKey` value for `thread_local!`
+    (@key $t:ty, $init:expr) => {{
+        #[inline]
+        fn __init() -> $t {
+            $init
+        }
+
+        unsafe {
+            use $crate::mem::needs_drop;
+            use $crate::thread::LocalKey;
+            use $crate::thread::local_impl::LazyStorage;
+
+            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)
+                    }
+                }
+            })
+        }
+    }},
+    ($(#[$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)*);
+    },
+}