1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

use super::abort_on_dtor_unwind;
use crate::cell::Cell;
use crate::marker::PhantomData;
use crate::ptr;
use crate::sys_common::thread_local_key::StaticKey as OsKey;

#[doc(hidden)]
#[allow_internal_unstable(thread_local_internals)]
#[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) => {
        $crate::thread::local_impl::thread_local_inner!(@key $t, { const INIT_EXPR: $t = $init; INIT_EXPR })
    },

    // used to generate the `LocalKey` value for `thread_local!`
    (@key $t:ty, $init:expr) => {
        {
            #[inline]
            fn __init() -> $t { $init }

            // `#[inline] does not work on windows-gnu due to linking errors around dllimports.
            // See https://github.com/rust-lang/rust/issues/109797.
            #[cfg_attr(not(windows), inline)]
            unsafe fn __getit(
                init: $crate::option::Option<&mut $crate::option::Option<$t>>,
            ) -> $crate::option::Option<&'static $t> {
                use $crate::thread::local_impl::Key;

                static __KEY: Key<$t> = Key::new();
                unsafe {
                    __KEY.get(init, __init)
                }
            }

            unsafe {
                $crate::thread::LocalKey::new(__getit)
            }
        }
    },
    ($(#[$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)*);
    },
}

/// Use a regular global static to store this key; the state provided will then be
/// thread-local.
#[allow(missing_debug_implementations)]
pub struct Key<T> {
    os: OsKey,
    marker: PhantomData<Cell<T>>,
}

unsafe impl<T> Sync for Key<T> {}

struct Value<T: 'static> {
    value: T,
    key: &'static Key<T>,
}

impl<T: 'static> Key<T> {
    #[rustc_const_unstable(feature = "thread_local_internals", issue = "none")]
    pub const fn new() -> Key<T> {
        Key { os: OsKey::new(Some(destroy_value::<T>)), marker: PhantomData }
    }

    /// Get the value associated with this key, initializating it if necessary.
    ///
    /// # Safety
    /// * the returned reference must not be used after recursive initialization
    /// or thread destruction occurs.
    pub unsafe fn get(
        &'static self,
        i: Option<&mut Option<T>>,
        f: impl FnOnce() -> T,
    ) -> Option<&'static T> {
        // SAFETY: (FIXME: get should actually be safe)
        let ptr = unsafe { self.os.get() as *mut Value<T> };
        if ptr.addr() > 1 {
            // SAFETY: the check ensured the pointer is safe (its destructor
            // is not running) + it is coming from a trusted source (self).
            unsafe { Some(&(*ptr).value) }
        } else {
            // SAFETY: At this point we are sure we have no value and so
            // initializing (or trying to) is safe.
            unsafe { self.try_initialize(ptr, i, f) }
        }
    }

    unsafe fn try_initialize(
        &'static self,
        ptr: *mut Value<T>,
        i: Option<&mut Option<T>>,
        f: impl FnOnce() -> T,
    ) -> Option<&'static T> {
        if ptr.addr() == 1 {
            // destructor is running
            return None;
        }

        let value = i.and_then(Option::take).unwrap_or_else(f);
        let ptr = Box::into_raw(Box::new(Value { value, key: self }));
        // SAFETY: (FIXME: get should actually be safe)
        let old = unsafe { self.os.get() as *mut Value<T> };
        // SAFETY: `ptr` is a correct pointer that can be destroyed by the key destructor.
        unsafe {
            self.os.set(ptr as *mut u8);
        }
        if !old.is_null() {
            // If the variable was recursively initialized, drop the old value.
            // SAFETY: We cannot be inside a `LocalKey::with` scope, as the
            // initializer has already returned and the next scope only starts
            // after we return the pointer. Therefore, there can be no references
            // to the old value.
            drop(unsafe { Box::from_raw(old) });
        }

        // SAFETY: We just created this value above.
        unsafe { Some(&(*ptr).value) }
    }
}

unsafe extern "C" fn destroy_value<T: 'static>(ptr: *mut u8) {
    // SAFETY:
    //
    // The OS TLS ensures that this key contains a null value when this
    // destructor starts to run. We set it back to a sentinel value of 1 to
    // ensure that any future calls to `get` for this thread will return
    // `None`.
    //
    // Note that to prevent an infinite loop we reset it back to null right
    // before we return from the destructor ourselves.
    abort_on_dtor_unwind(|| {
        let ptr = unsafe { Box::from_raw(ptr as *mut Value<T>) };
        let key = ptr.key;
        unsafe { key.os.set(ptr::without_provenance_mut(1)) };
        drop(ptr);
        unsafe { key.os.set(ptr::null_mut()) };
    });
}