1use crate::cell::{Cell, UnsafeCell};
2use crate::mem::MaybeUninit;
3use crate::ptr;
4use crate::sys::thread_local::{abort_on_dtor_unwind, destructors};
5
6pub unsafe trait DestroyedState: Sized + Copy {
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(t:ptr::from_ref(s).cast_mut().cast(), dtor:destroy::<T>);
18 }
19 }
20}
21
22#[derive(Copy, Clone)]
23enum State<D> {
24 Uninitialized,
25 Alive,
26 Destroyed(D),
27}
28
29#[allow(missing_debug_implementations)]
30pub struct Storage<T, D> {
31 state: Cell<State<D>>,
32 value: UnsafeCell<MaybeUninit<T>>,
33}
34
35impl<T, D> Storage<T, D>
36where
37 D: DestroyedState,
38{
39 pub const fn new() -> Storage<T, D> {
40 Storage {
41 state: Cell::new(State::Uninitialized),
42 value: UnsafeCell::new(MaybeUninit::uninit()),
43 }
44 }
45
46 /// Gets a pointer to the TLS value, potentially initializing it with the
47 /// provided parameters. If the TLS variable has been destroyed, a null
48 /// pointer is returned.
49 ///
50 /// The resulting pointer may not be used after reentrant inialialization
51 /// or thread destruction has occurred.
52 ///
53 /// # Safety
54 /// The `self` reference must remain valid until the TLS destructor is run.
55 #[inline]
56 pub unsafe fn get_or_init(&self, i: Option<&mut Option<T>>, f: impl FnOnce() -> T) -> *const T {
57 if let State::Alive = self.state.get() {
58 self.value.get().cast()
59 } else {
60 unsafe { self.get_or_init_slow(i, f) }
61 }
62 }
63
64 /// # Safety
65 /// The `self` reference must remain valid until the TLS destructor is run.
66 #[cold]
67 unsafe fn get_or_init_slow(
68 &self,
69 i: Option<&mut Option<T>>,
70 f: impl FnOnce() -> T,
71 ) -> *const T {
72 match self.state.get() {
73 State::Uninitialized => {}
74 State::Alive => return self.value.get().cast(),
75 State::Destroyed(_) => return ptr::null(),
76 }
77
78 let v = i.and_then(Option::take).unwrap_or_else(f);
79
80 // SAFETY: we cannot be inside a `LocalKey::with` scope, as the initializer
81 // has already returned and the next scope only starts after we return
82 // the pointer. Therefore, there can be no references to the old value,
83 // even if it was initialized. Thus because we are !Sync we have exclusive
84 // access to self.value and may replace it.
85 let mut old_value = unsafe { self.value.get().replace(MaybeUninit::new(v)) };
86 match self.state.replace(State::Alive) {
87 // If the variable is not being recursively initialized, register
88 // the destructor. This might be a noop if the value does not need
89 // destruction.
90 State::Uninitialized => D::register_dtor(self),
91
92 // Recursive initialization, we only need to drop the old value
93 // as we've already registered the destructor.
94 State::Alive => unsafe { old_value.assume_init_drop() },
95
96 State::Destroyed(_) => unreachable!(),
97 }
98
99 self.value.get().cast()
100 }
101}
102
103/// Transition an `Alive` TLS variable into the `Destroyed` state, dropping its
104/// value.
105///
106/// # Safety
107/// * Must only be called at thread destruction.
108/// * `ptr` must point to an instance of `Storage<T, ()>` and be valid for
109/// accessing that instance.
110unsafe extern "C" fn destroy<T>(ptr: *mut u8) {
111 // Print a nice abort message if a panic occurs.
112 abort_on_dtor_unwind(|| {
113 let storage: &Storage = unsafe { &*(ptr as *const Storage<T, ()>) };
114 if let State::Alive = storage.state.replace(val:State::Destroyed(())) {
115 // SAFETY: we ensured the state was Alive so the value was initialized.
116 // We also updated the state to Destroyed to prevent the destructor
117 // from accessing the thread-local variable, as this would violate
118 // the exclusive access provided by &mut T in Drop::drop.
119 unsafe { (*storage.value.get()).assume_init_drop() }
120 }
121 })
122}
123