1//! Thread-safe, non-blocking, "first one wins" flavor of `OnceCell`.
2//!
3//! If two threads race to initialize a type from the `race` module, they
4//! don't block, execute initialization function together, but only one of
5//! them stores the result.
6//!
7//! This module does not require `std` feature.
8//!
9//! # Atomic orderings
10//!
11//! All types in this module use `Acquire` and `Release`
12//! [atomic orderings](Ordering) for all their operations. While this is not
13//! strictly necessary for types other than `OnceBox`, it is useful for users as
14//! it allows them to be certain that after `get` or `get_or_init` returns on
15//! one thread, any side-effects caused by the setter thread prior to them
16//! calling `set` or `get_or_init` will be made visible to that thread; without
17//! it, it's possible for it to appear as if they haven't happened yet from the
18//! getter thread's perspective. This is an acceptable tradeoff to make since
19//! `Acquire` and `Release` have very little performance overhead on most
20//! architectures versus `Relaxed`.
21
22#[cfg(feature = "critical-section")]
23use atomic_polyfill as atomic;
24#[cfg(not(feature = "critical-section"))]
25use core::sync::atomic;
26
27use atomic::{AtomicPtr, AtomicUsize, Ordering};
28use core::cell::UnsafeCell;
29use core::marker::PhantomData;
30use core::num::NonZeroUsize;
31use core::ptr;
32
33/// A thread-safe cell which can be written to only once.
34#[derive(Default, Debug)]
35pub struct OnceNonZeroUsize {
36 inner: AtomicUsize,
37}
38
39impl OnceNonZeroUsize {
40 /// Creates a new empty cell.
41 #[inline]
42 pub const fn new() -> OnceNonZeroUsize {
43 OnceNonZeroUsize { inner: AtomicUsize::new(0) }
44 }
45
46 /// Gets the underlying value.
47 #[inline]
48 pub fn get(&self) -> Option<NonZeroUsize> {
49 let val = self.inner.load(Ordering::Acquire);
50 NonZeroUsize::new(val)
51 }
52
53 /// Sets the contents of this cell to `value`.
54 ///
55 /// Returns `Ok(())` if the cell was empty and `Err(())` if it was
56 /// full.
57 #[inline]
58 pub fn set(&self, value: NonZeroUsize) -> Result<(), ()> {
59 let exchange =
60 self.inner.compare_exchange(0, value.get(), Ordering::AcqRel, Ordering::Acquire);
61 match exchange {
62 Ok(_) => Ok(()),
63 Err(_) => Err(()),
64 }
65 }
66
67 /// Gets the contents of the cell, initializing it with `f` if the cell was
68 /// empty.
69 ///
70 /// If several threads concurrently run `get_or_init`, more than one `f` can
71 /// be called. However, all threads will return the same value, produced by
72 /// some `f`.
73 pub fn get_or_init<F>(&self, f: F) -> NonZeroUsize
74 where
75 F: FnOnce() -> NonZeroUsize,
76 {
77 enum Void {}
78 match self.get_or_try_init(|| Ok::<NonZeroUsize, Void>(f())) {
79 Ok(val) => val,
80 Err(void) => match void {},
81 }
82 }
83
84 /// Gets the contents of the cell, initializing it with `f` if
85 /// the cell was empty. If the cell was empty and `f` failed, an
86 /// error is returned.
87 ///
88 /// If several threads concurrently run `get_or_init`, more than one `f` can
89 /// be called. However, all threads will return the same value, produced by
90 /// some `f`.
91 pub fn get_or_try_init<F, E>(&self, f: F) -> Result<NonZeroUsize, E>
92 where
93 F: FnOnce() -> Result<NonZeroUsize, E>,
94 {
95 let val = self.inner.load(Ordering::Acquire);
96 let res = match NonZeroUsize::new(val) {
97 Some(it) => it,
98 None => {
99 let mut val = f()?.get();
100 let exchange =
101 self.inner.compare_exchange(0, val, Ordering::AcqRel, Ordering::Acquire);
102 if let Err(old) = exchange {
103 val = old;
104 }
105 unsafe { NonZeroUsize::new_unchecked(val) }
106 }
107 };
108 Ok(res)
109 }
110}
111
112/// A thread-safe cell which can be written to only once.
113#[derive(Default, Debug)]
114pub struct OnceBool {
115 inner: OnceNonZeroUsize,
116}
117
118impl OnceBool {
119 /// Creates a new empty cell.
120 #[inline]
121 pub const fn new() -> OnceBool {
122 OnceBool { inner: OnceNonZeroUsize::new() }
123 }
124
125 /// Gets the underlying value.
126 #[inline]
127 pub fn get(&self) -> Option<bool> {
128 self.inner.get().map(OnceBool::from_usize)
129 }
130
131 /// Sets the contents of this cell to `value`.
132 ///
133 /// Returns `Ok(())` if the cell was empty and `Err(())` if it was
134 /// full.
135 #[inline]
136 pub fn set(&self, value: bool) -> Result<(), ()> {
137 self.inner.set(OnceBool::to_usize(value))
138 }
139
140 /// Gets the contents of the cell, initializing it with `f` if the cell was
141 /// empty.
142 ///
143 /// If several threads concurrently run `get_or_init`, more than one `f` can
144 /// be called. However, all threads will return the same value, produced by
145 /// some `f`.
146 pub fn get_or_init<F>(&self, f: F) -> bool
147 where
148 F: FnOnce() -> bool,
149 {
150 OnceBool::from_usize(self.inner.get_or_init(|| OnceBool::to_usize(f())))
151 }
152
153 /// Gets the contents of the cell, initializing it with `f` if
154 /// the cell was empty. If the cell was empty and `f` failed, an
155 /// error is returned.
156 ///
157 /// If several threads concurrently run `get_or_init`, more than one `f` can
158 /// be called. However, all threads will return the same value, produced by
159 /// some `f`.
160 pub fn get_or_try_init<F, E>(&self, f: F) -> Result<bool, E>
161 where
162 F: FnOnce() -> Result<bool, E>,
163 {
164 self.inner.get_or_try_init(|| f().map(OnceBool::to_usize)).map(OnceBool::from_usize)
165 }
166
167 #[inline]
168 fn from_usize(value: NonZeroUsize) -> bool {
169 value.get() == 1
170 }
171
172 #[inline]
173 fn to_usize(value: bool) -> NonZeroUsize {
174 unsafe { NonZeroUsize::new_unchecked(if value { 1 } else { 2 }) }
175 }
176}
177
178/// A thread-safe cell which can be written to only once.
179pub struct OnceRef<'a, T> {
180 inner: AtomicPtr<T>,
181 ghost: PhantomData<UnsafeCell<&'a T>>,
182}
183
184// TODO: Replace UnsafeCell with SyncUnsafeCell once stabilized
185unsafe impl<'a, T: Sync> Sync for OnceRef<'a, T> {}
186
187impl<'a, T> core::fmt::Debug for OnceRef<'a, T> {
188 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
189 write!(f, "OnceRef({:?})", self.inner)
190 }
191}
192
193impl<'a, T> Default for OnceRef<'a, T> {
194 fn default() -> Self {
195 Self::new()
196 }
197}
198
199impl<'a, T> OnceRef<'a, T> {
200 /// Creates a new empty cell.
201 pub const fn new() -> OnceRef<'a, T> {
202 OnceRef { inner: AtomicPtr::new(ptr::null_mut()), ghost: PhantomData }
203 }
204
205 /// Gets a reference to the underlying value.
206 pub fn get(&self) -> Option<&'a T> {
207 let ptr = self.inner.load(Ordering::Acquire);
208 unsafe { ptr.as_ref() }
209 }
210
211 /// Sets the contents of this cell to `value`.
212 ///
213 /// Returns `Ok(())` if the cell was empty and `Err(value)` if it was
214 /// full.
215 pub fn set(&self, value: &'a T) -> Result<(), ()> {
216 let ptr = value as *const T as *mut T;
217 let exchange =
218 self.inner.compare_exchange(ptr::null_mut(), ptr, Ordering::AcqRel, Ordering::Acquire);
219 match exchange {
220 Ok(_) => Ok(()),
221 Err(_) => Err(()),
222 }
223 }
224
225 /// Gets the contents of the cell, initializing it with `f` if the cell was
226 /// empty.
227 ///
228 /// If several threads concurrently run `get_or_init`, more than one `f` can
229 /// be called. However, all threads will return the same value, produced by
230 /// some `f`.
231 pub fn get_or_init<F>(&self, f: F) -> &'a T
232 where
233 F: FnOnce() -> &'a T,
234 {
235 enum Void {}
236 match self.get_or_try_init(|| Ok::<&'a T, Void>(f())) {
237 Ok(val) => val,
238 Err(void) => match void {},
239 }
240 }
241
242 /// Gets the contents of the cell, initializing it with `f` if
243 /// the cell was empty. If the cell was empty and `f` failed, an
244 /// error is returned.
245 ///
246 /// If several threads concurrently run `get_or_init`, more than one `f` can
247 /// be called. However, all threads will return the same value, produced by
248 /// some `f`.
249 pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&'a T, E>
250 where
251 F: FnOnce() -> Result<&'a T, E>,
252 {
253 let mut ptr = self.inner.load(Ordering::Acquire);
254
255 if ptr.is_null() {
256 // TODO replace with `cast_mut` when MSRV reaches 1.65.0 (also in `set`)
257 ptr = f()? as *const T as *mut T;
258 let exchange = self.inner.compare_exchange(
259 ptr::null_mut(),
260 ptr,
261 Ordering::AcqRel,
262 Ordering::Acquire,
263 );
264 if let Err(old) = exchange {
265 ptr = old;
266 }
267 }
268
269 Ok(unsafe { &*ptr })
270 }
271
272 /// ```compile_fail
273 /// use once_cell::race::OnceRef;
274 ///
275 /// let mut l = OnceRef::new();
276 ///
277 /// {
278 /// let y = 2;
279 /// let mut r = OnceRef::new();
280 /// r.set(&y).unwrap();
281 /// core::mem::swap(&mut l, &mut r);
282 /// }
283 ///
284 /// // l now contains a dangling reference to y
285 /// eprintln!("uaf: {}", l.get().unwrap());
286 /// ```
287 fn _dummy() {}
288}
289
290#[cfg(feature = "alloc")]
291pub use self::once_box::OnceBox;
292
293#[cfg(feature = "alloc")]
294mod once_box {
295 use super::atomic::{AtomicPtr, Ordering};
296 use core::{marker::PhantomData, ptr};
297
298 use alloc::boxed::Box;
299
300 /// A thread-safe cell which can be written to only once.
301 pub struct OnceBox<T> {
302 inner: AtomicPtr<T>,
303 ghost: PhantomData<Option<Box<T>>>,
304 }
305
306 impl<T> core::fmt::Debug for OnceBox<T> {
307 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
308 write!(f, "OnceBox({:?})", self.inner.load(Ordering::Relaxed))
309 }
310 }
311
312 impl<T> Default for OnceBox<T> {
313 fn default() -> Self {
314 Self::new()
315 }
316 }
317
318 impl<T> Drop for OnceBox<T> {
319 fn drop(&mut self) {
320 let ptr = *self.inner.get_mut();
321 if !ptr.is_null() {
322 drop(unsafe { Box::from_raw(ptr) })
323 }
324 }
325 }
326
327 impl<T> OnceBox<T> {
328 /// Creates a new empty cell.
329 pub const fn new() -> OnceBox<T> {
330 OnceBox { inner: AtomicPtr::new(ptr::null_mut()), ghost: PhantomData }
331 }
332
333 /// Gets a reference to the underlying value.
334 pub fn get(&self) -> Option<&T> {
335 let ptr = self.inner.load(Ordering::Acquire);
336 if ptr.is_null() {
337 return None;
338 }
339 Some(unsafe { &*ptr })
340 }
341
342 /// Sets the contents of this cell to `value`.
343 ///
344 /// Returns `Ok(())` if the cell was empty and `Err(value)` if it was
345 /// full.
346 pub fn set(&self, value: Box<T>) -> Result<(), Box<T>> {
347 let ptr = Box::into_raw(value);
348 let exchange = self.inner.compare_exchange(
349 ptr::null_mut(),
350 ptr,
351 Ordering::AcqRel,
352 Ordering::Acquire,
353 );
354 if exchange.is_err() {
355 let value = unsafe { Box::from_raw(ptr) };
356 return Err(value);
357 }
358 Ok(())
359 }
360
361 /// Gets the contents of the cell, initializing it with `f` if the cell was
362 /// empty.
363 ///
364 /// If several threads concurrently run `get_or_init`, more than one `f` can
365 /// be called. However, all threads will return the same value, produced by
366 /// some `f`.
367 pub fn get_or_init<F>(&self, f: F) -> &T
368 where
369 F: FnOnce() -> Box<T>,
370 {
371 enum Void {}
372 match self.get_or_try_init(|| Ok::<Box<T>, Void>(f())) {
373 Ok(val) => val,
374 Err(void) => match void {},
375 }
376 }
377
378 /// Gets the contents of the cell, initializing it with `f` if
379 /// the cell was empty. If the cell was empty and `f` failed, an
380 /// error is returned.
381 ///
382 /// If several threads concurrently run `get_or_init`, more than one `f` can
383 /// be called. However, all threads will return the same value, produced by
384 /// some `f`.
385 pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&T, E>
386 where
387 F: FnOnce() -> Result<Box<T>, E>,
388 {
389 let mut ptr = self.inner.load(Ordering::Acquire);
390
391 if ptr.is_null() {
392 let val = f()?;
393 ptr = Box::into_raw(val);
394 let exchange = self.inner.compare_exchange(
395 ptr::null_mut(),
396 ptr,
397 Ordering::AcqRel,
398 Ordering::Acquire,
399 );
400 if let Err(old) = exchange {
401 drop(unsafe { Box::from_raw(ptr) });
402 ptr = old;
403 }
404 };
405 Ok(unsafe { &*ptr })
406 }
407 }
408
409 unsafe impl<T: Sync + Send> Sync for OnceBox<T> {}
410
411 /// ```compile_fail
412 /// struct S(*mut ());
413 /// unsafe impl Sync for S {}
414 ///
415 /// fn share<T: Sync>(_: &T) {}
416 /// share(&once_cell::race::OnceBox::<S>::new());
417 /// ```
418 fn _dummy() {}
419}
420