1 use crate::loom::thread::LocalKey;
2 
3 use std::cell::Cell;
4 use std::marker;
5 
6 /// Set a reference as a thread-local
7 macro_rules! scoped_thread_local {
8     ($(#[$attrs:meta])* $vis:vis static $name:ident: $ty:ty) => (
9         $(#[$attrs])*
10         $vis static $name: $crate::macros::scoped_tls::ScopedKey<$ty>
11             = $crate::macros::scoped_tls::ScopedKey {
12                 inner: {
13                     thread_local!(static FOO: ::std::cell::Cell<*const ()> = {
14                         std::cell::Cell::new(::std::ptr::null())
15                     });
16                     &FOO
17                 },
18                 _marker: ::std::marker::PhantomData,
19             };
20     )
21 }
22 
23 /// Type representing a thread local storage key corresponding to a reference
24 /// to the type parameter `T`.
25 pub(crate) struct ScopedKey<T> {
26     pub(crate) inner: &'static LocalKey<Cell<*const ()>>,
27     pub(crate) _marker: marker::PhantomData<T>,
28 }
29 
30 unsafe impl<T> Sync for ScopedKey<T> {}
31 
32 impl<T> ScopedKey<T> {
33     /// Inserts a value into this scoped thread local storage slot for a
34     /// duration of a closure.
set<F, R>(&'static self, t: &T, f: F) -> R where F: FnOnce() -> R,35     pub(crate) fn set<F, R>(&'static self, t: &T, f: F) -> R
36     where
37         F: FnOnce() -> R,
38     {
39         struct Reset {
40             key: &'static LocalKey<Cell<*const ()>>,
41             val: *const (),
42         }
43 
44         impl Drop for Reset {
45             fn drop(&mut self) {
46                 self.key.with(|c| c.set(self.val));
47             }
48         }
49 
50         let prev = self.inner.with(|c| {
51             let prev = c.get();
52             c.set(t as *const _ as *const ());
53             prev
54         });
55 
56         let _reset = Reset {
57             key: self.inner,
58             val: prev,
59         };
60 
61         f()
62     }
63 
64     /// Gets a value out of this scoped variable.
with<F, R>(&'static self, f: F) -> R where F: FnOnce(Option<&T>) -> R,65     pub(crate) fn with<F, R>(&'static self, f: F) -> R
66     where
67         F: FnOnce(Option<&T>) -> R,
68     {
69         let val = self.inner.with(|c| c.get());
70 
71         if val.is_null() {
72             f(None)
73         } else {
74             unsafe { f(Some(&*(val as *const T))) }
75         }
76     }
77 }
78