std/sys/thread_local/native/
lazy.rs

1use crate::cell::UnsafeCell;
2use crate::hint::unreachable_unchecked;
3use crate::ptr;
4use crate::sys::thread_local::{abort_on_dtor_unwind, destructors};
5
6pub unsafe trait DestroyedState: Sized {
7    fn register_dtor<T>(s: &Storage<T, Self>);
8}
9
10unsafe impl DestroyedState for ! {
11    fn register_dtor<T>(_: &Storage<T, !>) {}
12}
13
14unsafe impl DestroyedState for () {
15    fn register_dtor<T>(s: &Storage<T, ()>) {
16        unsafe {
17            destructors::register(ptr::from_ref(s).cast_mut().cast(), destroy::<T>);
18        }
19    }
20}
21
22enum State<T, D> {
23    Initial,
24    Alive(T),
25    Destroyed(D),
26}
27
28#[allow(missing_debug_implementations)]
29pub struct Storage<T, D> {
30    state: UnsafeCell<State<T, D>>,
31}
32
33impl<T, D> Storage<T, D>
34where
35    D: DestroyedState,
36{
37    pub const fn new() -> Storage<T, D> {
38        Storage { state: UnsafeCell::new(State::Initial) }
39    }
40
41    /// Gets a pointer to the TLS value, potentially initializing it with the
42    /// provided parameters. If the TLS variable has been destroyed, a null
43    /// pointer is returned.
44    ///
45    /// The resulting pointer may not be used after reentrant inialialization
46    /// or thread destruction has occurred.
47    ///
48    /// # Safety
49    /// The `self` reference must remain valid until the TLS destructor is run.
50    #[inline]
51    pub unsafe fn get_or_init(&self, i: Option<&mut Option<T>>, f: impl FnOnce() -> T) -> *const T {
52        let state = unsafe { &*self.state.get() };
53        match state {
54            State::Alive(v) => v,
55            State::Destroyed(_) => ptr::null(),
56            State::Initial => unsafe { self.initialize(i, f) },
57        }
58    }
59
60    #[cold]
61    unsafe fn initialize(&self, i: Option<&mut Option<T>>, f: impl FnOnce() -> T) -> *const T {
62        // Perform initialization
63
64        let v = i.and_then(Option::take).unwrap_or_else(f);
65
66        let old = unsafe { self.state.get().replace(State::Alive(v)) };
67        match old {
68            // If the variable is not being recursively initialized, register
69            // the destructor. This might be a noop if the value does not need
70            // destruction.
71            State::Initial => D::register_dtor(self),
72            // Else, drop the old value. This might be changed to a panic.
73            val => drop(val),
74        }
75
76        // SAFETY: the state was just set to `Alive`
77        unsafe {
78            let State::Alive(v) = &*self.state.get() else { unreachable_unchecked() };
79            v
80        }
81    }
82}
83
84/// Transition an `Alive` TLS variable into the `Destroyed` state, dropping its
85/// value.
86///
87/// # Safety
88/// * Must only be called at thread destruction.
89/// * `ptr` must point to an instance of `Storage<T, ()>` and be valid for
90///   accessing that instance.
91unsafe extern "C" fn destroy<T>(ptr: *mut u8) {
92    // Print a nice abort message if a panic occurs.
93    abort_on_dtor_unwind(|| {
94        let storage = unsafe { &*(ptr as *const Storage<T, ()>) };
95        // Update the state before running the destructor as it may attempt to
96        // access the variable.
97        let val = unsafe { storage.state.get().replace(State::Destroyed(())) };
98        drop(val);
99    })
100}
std/sys/thread_local/native/lazy.rs

std/sys/thread_local/native/
lazy.rs
