]>
Commit | Line | Data |
---|---|---|
e1599b0c XL |
1 | use std::{ |
2 | cell::UnsafeCell, | |
3dfed10e | 3 | hint, |
e1599b0c | 4 | panic::{RefUnwindSafe, UnwindSafe}, |
923072b8 | 5 | sync::atomic::{AtomicU8, Ordering}, |
e1599b0c XL |
6 | }; |
7 | ||
e1599b0c | 8 | pub(crate) struct OnceCell<T> { |
923072b8 | 9 | state: AtomicU8, |
3dfed10e | 10 | value: UnsafeCell<Option<T>>, |
e1599b0c XL |
11 | } |
12 | ||
923072b8 FG |
13 | const INCOMPLETE: u8 = 0x0; |
14 | const RUNNING: u8 = 0x1; | |
15 | const COMPLETE: u8 = 0x2; | |
16 | ||
e1599b0c XL |
17 | // Why do we need `T: Send`? |
18 | // Thread A creates a `OnceCell` and shares it with | |
19 | // scoped thread B, which fills the cell, which is | |
20 | // then destroyed by A. That is, destructor observes | |
21 | // a sent value. | |
22 | unsafe impl<T: Sync + Send> Sync for OnceCell<T> {} | |
23 | unsafe impl<T: Send> Send for OnceCell<T> {} | |
24 | ||
25 | impl<T: RefUnwindSafe + UnwindSafe> RefUnwindSafe for OnceCell<T> {} | |
26 | impl<T: UnwindSafe> UnwindSafe for OnceCell<T> {} | |
27 | ||
28 | impl<T> OnceCell<T> { | |
29 | pub(crate) const fn new() -> OnceCell<T> { | |
923072b8 FG |
30 | OnceCell { state: AtomicU8::new(INCOMPLETE), value: UnsafeCell::new(None) } |
31 | } | |
32 | ||
33 | pub(crate) const fn with_value(value: T) -> OnceCell<T> { | |
34 | OnceCell { state: AtomicU8::new(COMPLETE), value: UnsafeCell::new(Some(value)) } | |
e1599b0c XL |
35 | } |
36 | ||
37 | /// Safety: synchronizes with store to value via Release/Acquire. | |
38 | #[inline] | |
39 | pub(crate) fn is_initialized(&self) -> bool { | |
923072b8 | 40 | self.state.load(Ordering::Acquire) == COMPLETE |
e1599b0c XL |
41 | } |
42 | ||
43 | /// Safety: synchronizes with store to value via `is_initialized` or mutex | |
44 | /// lock/unlock, writes value only once because of the mutex. | |
45 | #[cold] | |
46 | pub(crate) fn initialize<F, E>(&self, f: F) -> Result<(), E> | |
47 | where | |
48 | F: FnOnce() -> Result<T, E>, | |
49 | { | |
6a06907d XL |
50 | let mut f = Some(f); |
51 | let mut res: Result<(), E> = Ok(()); | |
52 | let slot: *mut Option<T> = self.value.get(); | |
923072b8 | 53 | initialize_inner(&self.state, &mut || { |
e1599b0c XL |
54 | // We are calling user-supplied function and need to be careful. |
55 | // - if it returns Err, we unlock mutex and return without touching anything | |
56 | // - if it panics, we unlock mutex and propagate panic without touching anything | |
57 | // - if it calls `set` or `get_or_try_init` re-entrantly, we get a deadlock on | |
58 | // mutex, which is important for safety. We *could* detect this and panic, | |
59 | // but that is more complicated | |
60 | // - finally, if it returns Ok, we store the value and store the flag with | |
61 | // `Release`, which synchronizes with `Acquire`s. | |
923072b8 | 62 | let f = unsafe { crate::take_unchecked(&mut f) }; |
6a06907d XL |
63 | match f() { |
64 | Ok(value) => unsafe { | |
65 | // Safe b/c we have a unique access and no panic may happen | |
66 | // until the cell is marked as initialized. | |
67 | debug_assert!((*slot).is_none()); | |
68 | *slot = Some(value); | |
69 | true | |
70 | }, | |
71 | Err(err) => { | |
72 | res = Err(err); | |
73 | false | |
74 | } | |
75 | } | |
76 | }); | |
77 | res | |
e1599b0c | 78 | } |
f035d41b | 79 | |
923072b8 FG |
80 | #[cold] |
81 | pub(crate) fn wait(&self) { | |
82 | let key = &self.state as *const _ as usize; | |
83 | unsafe { | |
84 | parking_lot_core::park( | |
85 | key, | |
86 | || self.state.load(Ordering::Acquire) != COMPLETE, | |
87 | || (), | |
88 | |_, _| (), | |
89 | parking_lot_core::DEFAULT_PARK_TOKEN, | |
90 | None, | |
91 | ); | |
92 | } | |
93 | } | |
94 | ||
f035d41b XL |
95 | /// Get the reference to the underlying value, without checking if the cell |
96 | /// is initialized. | |
97 | /// | |
98 | /// # Safety | |
99 | /// | |
100 | /// Caller must ensure that the cell is in initialized state, and that | |
101 | /// the contents are acquired by (synchronized to) this thread. | |
102 | pub(crate) unsafe fn get_unchecked(&self) -> &T { | |
103 | debug_assert!(self.is_initialized()); | |
3dfed10e XL |
104 | let slot: &Option<T> = &*self.value.get(); |
105 | match slot { | |
106 | Some(value) => value, | |
107 | // This unsafe does improve performance, see `examples/bench`. | |
108 | None => { | |
109 | debug_assert!(false); | |
110 | hint::unreachable_unchecked() | |
111 | } | |
112 | } | |
f035d41b XL |
113 | } |
114 | ||
115 | /// Gets the mutable reference to the underlying value. | |
116 | /// Returns `None` if the cell is empty. | |
117 | pub(crate) fn get_mut(&mut self) -> Option<&mut T> { | |
3dfed10e XL |
118 | // Safe b/c we have an exclusive access |
119 | let slot: &mut Option<T> = unsafe { &mut *self.value.get() }; | |
120 | slot.as_mut() | |
f035d41b XL |
121 | } |
122 | ||
123 | /// Consumes this `OnceCell`, returning the wrapped value. | |
124 | /// Returns `None` if the cell was empty. | |
125 | pub(crate) fn into_inner(self) -> Option<T> { | |
3dfed10e | 126 | self.value.into_inner() |
e1599b0c XL |
127 | } |
128 | } | |
129 | ||
923072b8 FG |
130 | struct Guard<'a> { |
131 | state: &'a AtomicU8, | |
132 | new_state: u8, | |
133 | } | |
134 | ||
135 | impl<'a> Drop for Guard<'a> { | |
136 | fn drop(&mut self) { | |
137 | self.state.store(self.new_state, Ordering::Release); | |
138 | unsafe { | |
139 | let key = self.state as *const AtomicU8 as usize; | |
140 | parking_lot_core::unpark_all(key, parking_lot_core::DEFAULT_UNPARK_TOKEN); | |
141 | } | |
142 | } | |
143 | } | |
144 | ||
6a06907d XL |
145 | // Note: this is intentionally monomorphic |
146 | #[inline(never)] | |
923072b8 FG |
147 | fn initialize_inner(state: &AtomicU8, init: &mut dyn FnMut() -> bool) { |
148 | loop { | |
149 | let exchange = | |
150 | state.compare_exchange_weak(INCOMPLETE, RUNNING, Ordering::Acquire, Ordering::Acquire); | |
151 | match exchange { | |
152 | Ok(_) => { | |
153 | let mut guard = Guard { state, new_state: INCOMPLETE }; | |
154 | if init() { | |
155 | guard.new_state = COMPLETE; | |
156 | } | |
157 | return; | |
158 | } | |
159 | Err(COMPLETE) => return, | |
160 | Err(RUNNING) => unsafe { | |
161 | let key = state as *const AtomicU8 as usize; | |
162 | parking_lot_core::park( | |
163 | key, | |
164 | || state.load(Ordering::Relaxed) == RUNNING, | |
165 | || (), | |
166 | |_, _| (), | |
167 | parking_lot_core::DEFAULT_PARK_TOKEN, | |
168 | None, | |
169 | ); | |
170 | }, | |
171 | Err(_) => debug_assert!(false), | |
6a06907d XL |
172 | } |
173 | } | |
174 | } | |
175 | ||
e1599b0c | 176 | #[test] |
e1599b0c XL |
177 | fn test_size() { |
178 | use std::mem::size_of; | |
179 | ||
923072b8 | 180 | assert_eq!(size_of::<OnceCell<bool>>(), 1 * size_of::<bool>() + size_of::<u8>()); |
e1599b0c | 181 | } |