Support #[rustc_align_static] inside thread_local!

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

@@ -41,7 +41,7 @@ cfg_select! {
     }
     _ => {
         mod os;
-        pub use os::{Storage, thread_local_inner};
+        pub use os::{Storage, thread_local_inner, value_align};
         pub(crate) use os::{LocalPointer, local_pointer};
     }
 }

library/std/src/sys/thread_local/native/eager.rs

@@ -10,9 +10,11 @@ enum State {
 }
 
 #[allow(missing_debug_implementations)]
+#[repr(C)]
 pub struct Storage<T> {
-    state: Cell<State>,
+    // This field must be first, for correctness of `#[rustc_align_static]`
     val: UnsafeCell<T>,
+    state: Cell<State>,
 }
 
 impl<T> Storage<T> {

library/std/src/sys/thread_local/native/lazy.rs

@@ -27,9 +27,11 @@ enum State<D> {
 }
 
 #[allow(missing_debug_implementations)]
+#[repr(C)]
 pub struct Storage<T, D> {
-    state: Cell<State<D>>,
+    // This field must be first, for correctness of `#[rustc_align_static]`
     value: UnsafeCell<MaybeUninit<T>>,
+    state: Cell<State<D>>,
 }
 
 impl<T, D> Storage<T, D>

library/std/src/sys/thread_local/native/mod.rs

@@ -54,21 +54,23 @@ pub macro thread_local_inner {
     // test in `tests/thread.rs` if these types are renamed.
 
     // Used to generate the `LocalKey` value for const-initialized thread locals.
-    (@key $t:ty, const $init:expr) => {{
+    (@key $t:ty, $(#[$align_attr:meta])*, const $init:expr) => {{
         const __INIT: $t = $init;
 
         unsafe {
             $crate::thread::LocalKey::new(const {
                 if $crate::mem::needs_drop::<$t>() {
                     |_| {
                         #[thread_local]
+                        $(#[$align_attr])*
                         static VAL: $crate::thread::local_impl::EagerStorage<$t>
                             = $crate::thread::local_impl::EagerStorage::new(__INIT);
                         VAL.get()
                     }
                 } else {
                     |_| {
                         #[thread_local]
+                        $(#[$align_attr])*
                         static VAL: $t = __INIT;
                         &VAL
                     }
@@ -78,7 +80,7 @@ pub macro thread_local_inner {
     }},
 
     // used to generate the `LocalKey` value for `thread_local!`
-    (@key $t:ty, $init:expr) => {{
+    (@key $t:ty, $(#[$align_attr:meta])*, $init:expr) => {{
         #[inline]
         fn __init() -> $t {
             $init
@@ -89,13 +91,15 @@ pub macro thread_local_inner {
                 if $crate::mem::needs_drop::<$t>() {
                     |init| {
                         #[thread_local]
+                        $(#[$align_attr])*
                         static VAL: $crate::thread::local_impl::LazyStorage<$t, ()>
                             = $crate::thread::local_impl::LazyStorage::new();
                         VAL.get_or_init(init, __init)
                     }
                 } else {
                     |init| {
                         #[thread_local]
+                        $(#[$align_attr])*
                         static VAL: $crate::thread::local_impl::LazyStorage<$t, !>
                             = $crate::thread::local_impl::LazyStorage::new();
                         VAL.get_or_init(init, __init)
@@ -104,10 +108,6 @@ pub macro thread_local_inner {
             })
         }
     }},
-    ($(#[$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)*);
-    },
 }
 
 #[rustc_macro_transparency = "semitransparent"]

library/std/src/sys/thread_local/no_threads.rs

@@ -2,6 +2,7 @@
 //! thread locals and we can instead just use plain statics!
 
 use crate::cell::{Cell, UnsafeCell};
+use crate::mem::MaybeUninit;
 use crate::ptr;
 
 #[doc(hidden)]
@@ -11,12 +12,13 @@ use crate::ptr;
 #[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) => {{
+    (@key $t:ty, $(#[$align_attr:meta])*, const $init:expr) => {{
         const __INIT: $t = $init;
 
         // NOTE: Please update the shadowing test in `tests/thread.rs` if these types are renamed.
         unsafe {
             $crate::thread::LocalKey::new(|_| {
+                $(#[$align_attr])*
                 static VAL: $crate::thread::local_impl::EagerStorage<$t> =
                     $crate::thread::local_impl::EagerStorage { value: __INIT };
                 &VAL.value
@@ -25,42 +27,50 @@ pub macro thread_local_inner {
     }},
 
     // used to generate the `LocalKey` value for `thread_local!`
-    (@key $t:ty, $init:expr) => {{
+    (@key $t:ty, $(#[$align_attr:meta])*, $init:expr) => {{
         #[inline]
         fn __init() -> $t { $init }
 
         unsafe {
-            use $crate::thread::LocalKey;
-            use $crate::thread::local_impl::LazyStorage;
-
-            LocalKey::new(|init| {
-                static VAL: LazyStorage<$t> = LazyStorage::new();
+            $crate::thread::LocalKey::new(|init| {
+                $(#[$align_attr])*
+                static VAL: $crate::thread::local_impl::LazyStorage<$t> = $crate::thread::local_impl::LazyStorage::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)*);
-    },
 }
 
 #[allow(missing_debug_implementations)]
+#[repr(transparent)] // Required for correctness of `#[rustc_align_static]`
 pub struct EagerStorage<T> {
     pub value: T,
 }
 
 // SAFETY: the target doesn't have threads.
 unsafe impl<T> Sync for EagerStorage<T> {}
 
+#[derive(Clone, Copy, PartialEq, Eq)]
+enum State {
+    Initial,
+    Alive,
+    Destroying,
+}
+
 #[allow(missing_debug_implementations)]
+#[repr(C)]
 pub struct LazyStorage<T> {
-    value: UnsafeCell<Option<T>>,
+    // This field must be first, for correctness of `#[rustc_align_static]`
+    value: UnsafeCell<MaybeUninit<T>>,
+    state: Cell<State>,
 }
 
 impl<T> LazyStorage<T> {
     pub const fn new() -> LazyStorage<T> {
-        LazyStorage { value: UnsafeCell::new(None) }
+        LazyStorage {
+            value: UnsafeCell::new(MaybeUninit::uninit()),
+            state: Cell::new(State::Initial),
+        }
     }
 
     /// Gets a pointer to the TLS value, potentially initializing it with the
@@ -70,24 +80,39 @@ impl<T> LazyStorage<T> {
     /// has occurred.
     #[inline]
     pub fn get(&'static self, i: Option<&mut Option<T>>, f: impl FnOnce() -> T) -> *const T {
-        let value = unsafe { &*self.value.get() };
-        match value {
-            Some(v) => v,
-            None => self.initialize(i, f),
+        if self.state.get() == State::Alive {
+            self.value.get() as *const T
+        } else {
+            self.initialize(i, f)
         }
     }
 
     #[cold]
     fn initialize(&'static self, i: Option<&mut Option<T>>, f: impl FnOnce() -> T) -> *const 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.
+
+        // Destroy the old value if it is initialized
         // FIXME(#110897): maybe panic on recursive initialization.
+        if self.state.get() == State::Alive {
+            self.state.set(State::Destroying);
+            // Safety: we check for no initialization during drop below
+            unsafe {
+                ptr::drop_in_place(self.value.get() as *mut T);
+            }
+            self.state.set(State::Initial);
+        }
+
+        // Guard against initialization during drop
+        if self.state.get() == State::Destroying {
+            panic!("Attempted to initialize thread-local while it is being dropped");
+        }
+
         unsafe {
-            self.value.get().replace(Some(value));
+            self.value.get().write(MaybeUninit::new(value));
         }
-        // SAFETY: we just set this to `Some`.
-        unsafe { (*self.value.get()).as_ref().unwrap_unchecked() }
+        self.state.set(State::Alive);
+
+        self.value.get() as *const T
     }
 }