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// The "atomic orderings" section of the documentation above promises
23// "happens-before" semantics. This drives the choice of orderings in the uses
24// of `compare_exchange` below. On success, the value was zero/null, so there
25// was nothing to acquire (there is never any `Ordering::Release` store of 0).
26// On failure, the value was nonzero, so it was initialized previously (perhaps
27// on another thread) using `Ordering::Release`, so we must use
28// `Ordering::Acquire` to ensure that store "happens-before" this load.
29
30#[cfg(not(feature = "portable-atomic"))]
31use core::sync::atomic;
32#[cfg(feature = "portable-atomic")]
33use portable_atomic as atomic;
34
35use atomic::{AtomicPtr, AtomicUsize, Ordering};
36use core::cell::UnsafeCell;
37use core::marker::PhantomData;
38use core::num::NonZeroUsize;
39use core::ptr;
40
41/// A thread-safe cell which can be written to only once.
42#[derive(Default, Debug)]
43pub struct OnceNonZeroUsize {
44 inner: AtomicUsize,
45}
46
47impl OnceNonZeroUsize {
48 /// Creates a new empty cell.
49 #[inline]
50 pub const fn new() -> Self {
51 Self { inner: AtomicUsize::new(0) }
52 }
53
54 /// Gets the underlying value.
55 #[inline]
56 pub fn get(&self) -> Option<NonZeroUsize> {
57 let val = self.inner.load(Ordering::Acquire);
58 NonZeroUsize::new(val)
59 }
60
61 /// Get the reference to the underlying value, without checking if the cell
62 /// is initialized.
63 ///
64 /// # Safety
65 ///
66 /// Caller must ensure that the cell is in initialized state, and that
67 /// the contents are acquired by (synchronized to) this thread.
68 pub unsafe fn get_unchecked(&self) -> NonZeroUsize {
69 #[inline(always)]
70 fn as_const_ptr(r: &AtomicUsize) -> *const usize {
71 use core::mem::align_of;
72
73 let p: *const AtomicUsize = r;
74 // SAFETY: "This type has the same size and bit validity as
75 // the underlying integer type, usize. However, the alignment of
76 // this type is always equal to its size, even on targets where
77 // usize has a lesser alignment."
78 const _ALIGNMENT_COMPATIBLE: () =
79 assert!(align_of::<AtomicUsize>() % align_of::<usize>() == 0);
80 p.cast::<usize>()
81 }
82
83 // TODO(MSRV-1.70): Use `AtomicUsize::as_ptr().cast_const()`
84 // See https://github.com/rust-lang/rust/issues/138246.
85 let p = as_const_ptr(&self.inner);
86
87 // SAFETY: The caller is responsible for ensuring that the value
88 // was initialized and that the contents have been acquired by
89 // this thread. Assuming that, we can assume there will be no
90 // conflicting writes to the value since the value will never
91 // change once initialized. This relies on the statement in
92 // https://doc.rust-lang.org/1.83.0/core/sync/atomic/ that "(A
93 // `compare_exchange` or `compare_exchange_weak` that does not
94 // succeed is not considered a write."
95 let val = unsafe { p.read() };
96
97 // SAFETY: The caller is responsible for ensuring the value is
98 // initialized and thus not zero.
99 unsafe { NonZeroUsize::new_unchecked(val) }
100 }
101
102 /// Sets the contents of this cell to `value`.
103 ///
104 /// Returns `Ok(())` if the cell was empty and `Err(())` if it was
105 /// full.
106 #[inline]
107 pub fn set(&self, value: NonZeroUsize) -> Result<(), ()> {
108 match self.compare_exchange(value) {
109 Ok(_) => Ok(()),
110 Err(_) => Err(()),
111 }
112 }
113
114 /// Gets the contents of the cell, initializing it with `f` if the cell was
115 /// empty.
116 ///
117 /// If several threads concurrently run `get_or_init`, more than one `f` can
118 /// be called. However, all threads will return the same value, produced by
119 /// some `f`.
120 pub fn get_or_init<F>(&self, f: F) -> NonZeroUsize
121 where
122 F: FnOnce() -> NonZeroUsize,
123 {
124 enum Void {}
125 match self.get_or_try_init(|| Ok::<NonZeroUsize, Void>(f())) {
126 Ok(val) => val,
127 Err(void) => match void {},
128 }
129 }
130
131 /// Gets the contents of the cell, initializing it with `f` if
132 /// the cell was empty. If the cell was empty and `f` failed, an
133 /// error is returned.
134 ///
135 /// If several threads concurrently run `get_or_init`, more than one `f` can
136 /// be called. However, all threads will return the same value, produced by
137 /// some `f`.
138 pub fn get_or_try_init<F, E>(&self, f: F) -> Result<NonZeroUsize, E>
139 where
140 F: FnOnce() -> Result<NonZeroUsize, E>,
141 {
142 match self.get() {
143 Some(it) => Ok(it),
144 None => self.init(f),
145 }
146 }
147
148 #[cold]
149 #[inline(never)]
150 fn init<E>(&self, f: impl FnOnce() -> Result<NonZeroUsize, E>) -> Result<NonZeroUsize, E> {
151 let nz = f()?;
152 let mut val = nz.get();
153 if let Err(old) = self.compare_exchange(nz) {
154 val = old;
155 }
156 Ok(unsafe { NonZeroUsize::new_unchecked(val) })
157 }
158
159 #[inline(always)]
160 fn compare_exchange(&self, val: NonZeroUsize) -> Result<usize, usize> {
161 self.inner.compare_exchange(0, val.get(), Ordering::Release, Ordering::Acquire)
162 }
163}
164
165/// A thread-safe cell which can be written to only once.
166#[derive(Default, Debug)]
167pub struct OnceBool {
168 inner: OnceNonZeroUsize,
169}
170
171impl OnceBool {
172 /// Creates a new empty cell.
173 #[inline]
174 pub const fn new() -> Self {
175 Self { inner: OnceNonZeroUsize::new() }
176 }
177
178 /// Gets the underlying value.
179 #[inline]
180 pub fn get(&self) -> Option<bool> {
181 self.inner.get().map(Self::from_usize)
182 }
183
184 /// Sets the contents of this cell to `value`.
185 ///
186 /// Returns `Ok(())` if the cell was empty and `Err(())` if it was
187 /// full.
188 #[inline]
189 pub fn set(&self, value: bool) -> Result<(), ()> {
190 self.inner.set(Self::to_usize(value))
191 }
192
193 /// Gets the contents of the cell, initializing it with `f` if the cell was
194 /// empty.
195 ///
196 /// If several threads concurrently run `get_or_init`, more than one `f` can
197 /// be called. However, all threads will return the same value, produced by
198 /// some `f`.
199 pub fn get_or_init<F>(&self, f: F) -> bool
200 where
201 F: FnOnce() -> bool,
202 {
203 Self::from_usize(self.inner.get_or_init(|| Self::to_usize(f())))
204 }
205
206 /// Gets the contents of the cell, initializing it with `f` if
207 /// the cell was empty. If the cell was empty and `f` failed, an
208 /// error is returned.
209 ///
210 /// If several threads concurrently run `get_or_init`, more than one `f` can
211 /// be called. However, all threads will return the same value, produced by
212 /// some `f`.
213 pub fn get_or_try_init<F, E>(&self, f: F) -> Result<bool, E>
214 where
215 F: FnOnce() -> Result<bool, E>,
216 {
217 self.inner.get_or_try_init(|| f().map(Self::to_usize)).map(Self::from_usize)
218 }
219
220 #[inline]
221 fn from_usize(value: NonZeroUsize) -> bool {
222 value.get() == 1
223 }
224
225 #[inline]
226 fn to_usize(value: bool) -> NonZeroUsize {
227 unsafe { NonZeroUsize::new_unchecked(if value { 1 } else { 2 }) }
228 }
229}
230
231/// A thread-safe cell which can be written to only once.
232pub struct OnceRef<'a, T> {
233 inner: AtomicPtr<T>,
234 ghost: PhantomData<UnsafeCell<&'a T>>,
235}
236
237// TODO: Replace UnsafeCell with SyncUnsafeCell once stabilized
238unsafe impl<'a, T: Sync> Sync for OnceRef<'a, T> {}
239
240impl<'a, T> core::fmt::Debug for OnceRef<'a, T> {
241 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
242 write!(f, "OnceRef({:?})", self.inner)
243 }
244}
245
246impl<'a, T> Default for OnceRef<'a, T> {
247 fn default() -> Self {
248 Self::new()
249 }
250}
251
252impl<'a, T> OnceRef<'a, T> {
253 /// Creates a new empty cell.
254 pub const fn new() -> Self {
255 Self { inner: AtomicPtr::new(ptr::null_mut()), ghost: PhantomData }
256 }
257
258 /// Gets a reference to the underlying value.
259 pub fn get(&self) -> Option<&'a T> {
260 let ptr = self.inner.load(Ordering::Acquire);
261 unsafe { ptr.as_ref() }
262 }
263
264 /// Sets the contents of this cell to `value`.
265 ///
266 /// Returns `Ok(())` if the cell was empty and `Err(value)` if it was
267 /// full.
268 pub fn set(&self, value: &'a T) -> Result<(), ()> {
269 match self.compare_exchange(value) {
270 Ok(_) => Ok(()),
271 Err(_) => Err(()),
272 }
273 }
274
275 /// Gets the contents of the cell, initializing it with `f` if the cell was
276 /// empty.
277 ///
278 /// If several threads concurrently run `get_or_init`, more than one `f` can
279 /// be called. However, all threads will return the same value, produced by
280 /// some `f`.
281 pub fn get_or_init<F>(&self, f: F) -> &'a T
282 where
283 F: FnOnce() -> &'a T,
284 {
285 enum Void {}
286 match self.get_or_try_init(|| Ok::<&'a T, Void>(f())) {
287 Ok(val) => val,
288 Err(void) => match void {},
289 }
290 }
291
292 /// Gets the contents of the cell, initializing it with `f` if
293 /// the cell was empty. If the cell was empty and `f` failed, an
294 /// error is returned.
295 ///
296 /// If several threads concurrently run `get_or_init`, more than one `f` can
297 /// be called. However, all threads will return the same value, produced by
298 /// some `f`.
299 pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&'a T, E>
300 where
301 F: FnOnce() -> Result<&'a T, E>,
302 {
303 match self.get() {
304 Some(val) => Ok(val),
305 None => self.init(f),
306 }
307 }
308
309 #[cold]
310 #[inline(never)]
311 fn init<E>(&self, f: impl FnOnce() -> Result<&'a T, E>) -> Result<&'a T, E> {
312 let mut value: &'a T = f()?;
313 if let Err(old) = self.compare_exchange(value) {
314 value = unsafe { &*old };
315 }
316 Ok(value)
317 }
318
319 #[inline(always)]
320 fn compare_exchange(&self, value: &'a T) -> Result<(), *const T> {
321 self.inner
322 .compare_exchange(
323 ptr::null_mut(),
324 <*const T>::cast_mut(value),
325 Ordering::Release,
326 Ordering::Acquire,
327 )
328 .map(|_: *mut T| ())
329 .map_err(<*mut T>::cast_const)
330 }
331
332 /// ```compile_fail
333 /// use once_cell::race::OnceRef;
334 ///
335 /// let mut l = OnceRef::new();
336 ///
337 /// {
338 /// let y = 2;
339 /// let mut r = OnceRef::new();
340 /// r.set(&y).unwrap();
341 /// core::mem::swap(&mut l, &mut r);
342 /// }
343 ///
344 /// // l now contains a dangling reference to y
345 /// eprintln!("uaf: {}", l.get().unwrap());
346 /// ```
347 fn _dummy() {}
348}
349
350#[cfg(feature = "alloc")]
351pub use self::once_box::OnceBox;
352
353#[cfg(feature = "alloc")]
354mod once_box {
355 use super::atomic::{AtomicPtr, Ordering};
356 use core::{marker::PhantomData, ptr};
357
358 use alloc::boxed::Box;
359
360 /// A thread-safe cell which can be written to only once.
361 pub struct OnceBox<T> {
362 inner: AtomicPtr<T>,
363 ghost: PhantomData<Option<Box<T>>>,
364 }
365
366 impl<T> core::fmt::Debug for OnceBox<T> {
367 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
368 write!(f, "OnceBox({:?})", self.inner.load(Ordering::Relaxed))
369 }
370 }
371
372 impl<T> Default for OnceBox<T> {
373 fn default() -> Self {
374 Self::new()
375 }
376 }
377
378 impl<T> Drop for OnceBox<T> {
379 fn drop(&mut self) {
380 let ptr = *self.inner.get_mut();
381 if !ptr.is_null() {
382 drop(unsafe { Box::from_raw(ptr) })
383 }
384 }
385 }
386
387 impl<T> OnceBox<T> {
388 /// Creates a new empty cell.
389 pub const fn new() -> Self {
390 Self { inner: AtomicPtr::new(ptr::null_mut()), ghost: PhantomData }
391 }
392
393 /// Creates a new cell with the given value.
394 pub fn with_value(value: Box<T>) -> Self {
395 Self { inner: AtomicPtr::new(Box::into_raw(value)), ghost: PhantomData }
396 }
397
398 /// Gets a reference to the underlying value.
399 pub fn get(&self) -> Option<&T> {
400 let ptr = self.inner.load(Ordering::Acquire);
401 if ptr.is_null() {
402 return None;
403 }
404 Some(unsafe { &*ptr })
405 }
406
407 /// Sets the contents of this cell to `value`.
408 ///
409 /// Returns `Ok(())` if the cell was empty and `Err(value)` if it was
410 /// full.
411 pub fn set(&self, value: Box<T>) -> Result<(), Box<T>> {
412 let ptr = Box::into_raw(value);
413 let exchange = self.inner.compare_exchange(
414 ptr::null_mut(),
415 ptr,
416 Ordering::Release,
417 Ordering::Acquire,
418 );
419 if exchange.is_err() {
420 let value = unsafe { Box::from_raw(ptr) };
421 return Err(value);
422 }
423 Ok(())
424 }
425
426 /// Gets the contents of the cell, initializing it with `f` if the cell was
427 /// empty.
428 ///
429 /// If several threads concurrently run `get_or_init`, more than one `f` can
430 /// be called. However, all threads will return the same value, produced by
431 /// some `f`.
432 pub fn get_or_init<F>(&self, f: F) -> &T
433 where
434 F: FnOnce() -> Box<T>,
435 {
436 enum Void {}
437 match self.get_or_try_init(|| Ok::<Box<T>, Void>(f())) {
438 Ok(val) => val,
439 Err(void) => match void {},
440 }
441 }
442
443 /// Gets the contents of the cell, initializing it with `f` if
444 /// the cell was empty. If the cell was empty and `f` failed, an
445 /// error is returned.
446 ///
447 /// If several threads concurrently run `get_or_init`, more than one `f` can
448 /// be called. However, all threads will return the same value, produced by
449 /// some `f`.
450 pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&T, E>
451 where
452 F: FnOnce() -> Result<Box<T>, E>,
453 {
454 match self.get() {
455 Some(val) => Ok(val),
456 None => self.init(f)
457 }
458 }
459
460 #[cold]
461 #[inline(never)]
462 fn init<E>(&self, f: impl FnOnce() -> Result<Box<T>, E>) -> Result<&T, E> {
463 let val = f()?;
464 let mut ptr = Box::into_raw(val);
465 let exchange = self.inner.compare_exchange(
466 ptr::null_mut(),
467 ptr,
468 Ordering::Release,
469 Ordering::Acquire,
470 );
471 if let Err(old) = exchange {
472 drop(unsafe { Box::from_raw(ptr) });
473 ptr = old;
474 }
475 Ok(unsafe { &*ptr })
476 }
477 }
478
479 unsafe impl<T: Sync + Send> Sync for OnceBox<T> {}
480
481 impl<T: Clone> Clone for OnceBox<T> {
482 fn clone(&self) -> Self {
483 match self.get() {
484 Some(value) => OnceBox::with_value(Box::new(value.clone())),
485 None => OnceBox::new(),
486 }
487 }
488 }
489
490 /// ```compile_fail
491 /// struct S(*mut ());
492 /// unsafe impl Sync for S {}
493 ///
494 /// fn share<T: Sync>(_: &T) {}
495 /// share(&once_cell::race::OnceBox::<S>::new());
496 /// ```
497 fn _dummy() {}
498}
499