os impl: use const generic for alignment

Author: Jules-Bertholet

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/os.rs

@@ -30,30 +30,30 @@ pub macro thread_local_inner {
         // in `tests/thread.rs` if these types are renamed.
         unsafe {
             $crate::thread::LocalKey::new(|init| {
-                static VAL: $crate::thread::local_impl::Storage<$t>
+                const ALIGN: usize = $crate::thread::local_impl::thread_local_inner!(@align $t, $(#[$($align_attr)*])*);
+                static VAL: $crate::thread::local_impl::Storage<$t, ALIGN>
                     = $crate::thread::local_impl::Storage::new();
-                VAL.get($crate::thread::local_impl::thread_local_inner!(@align $(#[$($align_attr)*])*), init, __init)
+                VAL.get(init, __init)
             })
         }
     }},
 
     // Handle `rustc_align_static` attributes,
     // by translating them into an argumemt to pass to `Storage::get`:
 
-    // fast path for when there are none
-    (@align) => (1),
-
     // `rustc_align_static` attributes are present,
     // translate them into a `const` block that computes the alignment
-    (@align $(#[$($attr:tt)*])+) => {
+    (@align $t:ty, $(#[$($attr:tt)*])*) => {
         const {
             // Ensure that attributes have valid syntax
             // and that the proper feature gate is enabled
-            $(#[$($attr)*])+
+            $(#[$($attr)*])*
+            #[allow(unused)]
             static DUMMY: () = ();
 
-            let mut final_align = 1_usize;
-            $($crate::thread::local_impl::thread_local_inner!(@align_single final_align, $($attr)*);)+
+            #[allow(unused_mut)]
+            let mut final_align = $crate::thread::local_impl::value_align::<$t>();
+            $($crate::thread::local_impl::thread_local_inner!(@align_single final_align, $($attr)*);)*
             final_align
         }
     },
@@ -107,26 +107,30 @@ pub macro thread_local_inner {
 
 /// Use a regular global static to store this key; the state provided will then be
 /// thread-local.
+/// INVARIANT: ALIGN must be a valid alignment, and no less than `value_align::<T>`.
 #[allow(missing_debug_implementations)]
-pub struct Storage<T> {
+pub struct Storage<T, const ALIGN: usize> {
     key: LazyKey,
     marker: PhantomData<Cell<T>>,
 }
 
-unsafe impl<T> Sync for Storage<T> {}
+unsafe impl<T, const ALIGN: usize> Sync for Storage<T, ALIGN> {}
 
 #[repr(C)]
 struct Value<T: 'static> {
     // This field must be first, for correctness of `#[rustc_align_static]`
     value: T,
-    align: usize,
     // INVARIANT: if this value is stored under a TLS key, `key` must be that `key`.
     key: Key,
 }
 
-impl<T: 'static> Storage<T> {
-    pub const fn new() -> Storage<T> {
-        Storage { key: LazyKey::new(Some(destroy_value::<T>)), marker: PhantomData }
+pub const fn value_align<T: 'static>() -> usize {
+    crate::mem::align_of::<Value<T>>()
+}
+
+impl<T: 'static, const ALIGN: usize> Storage<T, ALIGN> {
+    pub const fn new() -> Storage<T, ALIGN> {
+        Storage { key: LazyKey::new(Some(destroy_value::<T, ALIGN>)), marker: PhantomData }
     }
 
     /// Gets a pointer to the TLS value, potentially initializing it with the
@@ -135,12 +139,7 @@ impl<T: 'static> Storage<T> {
     ///
     /// The resulting pointer may not be used after reentrant inialialization
     /// or thread destruction has occurred.
-    pub fn get(
-        &'static self,
-        align: usize,
-        i: Option<&mut Option<T>>,
-        f: impl FnOnce() -> T,
-    ) -> *const T {
+    pub fn get(&'static self, i: Option<&mut Option<T>>, f: impl FnOnce() -> T) -> *const T {
         let key = self.key.force();
         let ptr = unsafe { get(key) as *mut Value<T> };
         if ptr.addr() > 1 {
@@ -149,7 +148,7 @@ impl<T: 'static> Storage<T> {
             unsafe { &(*ptr).value }
         } else {
             // SAFETY: trivially correct.
-            unsafe { Self::try_initialize(key, align, ptr, i, f) }
+            unsafe { Self::try_initialize(key, ptr, i, f) }
         }
     }
 
@@ -158,7 +157,6 @@ impl<T: 'static> Storage<T> {
     /// * `ptr` must be the current value associated with `key`.
     unsafe fn try_initialize(
         key: Key,
-        align: usize,
         ptr: *mut Value<T>,
         i: Option<&mut Option<T>>,
         f: impl FnOnce() -> T,
@@ -169,17 +167,13 @@ impl<T: 'static> Storage<T> {
         }
 
         // Manually allocate with the requested alignment
-        let layout = Layout::new::<Value<T>>().align_to(align).unwrap();
+        let layout = Layout::new::<Value<T>>().align_to(ALIGN).unwrap();
         let ptr: *mut Value<T> = (unsafe { crate::alloc::alloc(layout) }).cast();
         if ptr.is_null() {
             crate::alloc::handle_alloc_error(layout);
         }
         unsafe {
-            ptr.write(Value {
-                value: i.and_then(Option::take).unwrap_or_else(f),
-                align: layout.align(),
-                key,
-            });
+            ptr.write(Value { value: i.and_then(Option::take).unwrap_or_else(f), key });
         }
 
         // SAFETY:
@@ -209,7 +203,7 @@ impl<T: 'static> Storage<T> {
     }
 }
 
-unsafe extern "C" fn destroy_value<T: 'static>(ptr: *mut u8) {
+unsafe extern "C" fn destroy_value<T: 'static, const ALIGN: usize>(ptr: *mut u8) {
     // SAFETY:
     //
     // The OS TLS ensures that this key contains a null value when this
@@ -223,14 +217,13 @@ unsafe extern "C" fn destroy_value<T: 'static>(ptr: *mut u8) {
         let value_ptr: *mut Value<T> = ptr.cast();
         unsafe {
             let key = (*value_ptr).key;
-            let align = (*value_ptr).align;
 
             // SAFETY: `key` is the TLS key `ptr` was stored under.
             set(key, ptr::without_provenance_mut(1));
 
             // drop and deallocate the value
             let layout =
-                Layout::from_size_align_unchecked(crate::mem::size_of::<Value<T>>(), align);
+                Layout::from_size_align_unchecked(crate::mem::size_of::<Value<T>>(), ALIGN);
             value_ptr.drop_in_place();
             crate::alloc::dealloc(ptr, layout);