vendor/once_cell/src/race.rs

   1 //! "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 use core::{
  10     num::NonZeroUsize,
  11     sync::atomic::{AtomicUsize, Ordering},
  12 };
  13
  14 /// A thread-safe cell which can be written to only once.
  15 #[derive(Default, Debug)]
  16 pub struct OnceNonZeroUsize {
  17     inner: AtomicUsize,
  18 }
  19
  20 impl OnceNonZeroUsize {
  21     /// Creates a new empty cell.
  22     #[inline]
  23     pub const fn new() -> OnceNonZeroUsize {
  24         OnceNonZeroUsize { inner: AtomicUsize::new(0) }
  25     }
  26
  27     /// Gets the underlying value.
  28     #[inline]
  29     pub fn get(&self) -> Option<NonZeroUsize> {
  30         let val = self.inner.load(Ordering::Acquire);
  31         NonZeroUsize::new(val)
  32     }
  33
  34     /// Sets the contents of this cell to `value`.
  35     ///
  36     /// Returns `Ok(())` if the cell was empty and `Err(())` if it was
  37     /// full.
  38     #[inline]
  39     pub fn set(&self, value: NonZeroUsize) -> Result<(), ()> {
  40         let exchange =
  41             self.inner.compare_exchange(0, value.get(), Ordering::AcqRel, Ordering::Acquire);
  42         match exchange {
  43             Ok(_) => Ok(()),
  44             Err(_) => Err(()),
  45         }
  46     }
  47
  48     /// Gets the contents of the cell, initializing it with `f` if the cell was
  49     /// empty.
  50     ///
  51     /// If several threads concurrently run `get_or_init`, more than one `f` can
  52     /// be called. However, all threads will return the same value, produced by
  53     /// some `f`.
  54     pub fn get_or_init<F>(&self, f: F) -> NonZeroUsize
  55     where
  56         F: FnOnce() -> NonZeroUsize,
  57     {
  58         enum Void {}
  59         match self.get_or_try_init(|| Ok::<NonZeroUsize, Void>(f())) {
  60             Ok(val) => val,
  61             Err(void) => match void {},
  62         }
  63     }
  64
  65     /// Gets the contents of the cell, initializing it with `f` if
  66     /// the cell was empty. If the cell was empty and `f` failed, an
  67     /// error is returned.
  68     ///
  69     /// If several threads concurrently run `get_or_init`, more than one `f` can
  70     /// be called. However, all threads will return the same value, produced by
  71     /// some `f`.
  72     pub fn get_or_try_init<F, E>(&self, f: F) -> Result<NonZeroUsize, E>
  73     where
  74         F: FnOnce() -> Result<NonZeroUsize, E>,
  75     {
  76         let val = self.inner.load(Ordering::Acquire);
  77         let res = match NonZeroUsize::new(val) {
  78             Some(it) => it,
  79             None => {
  80                 let mut val = f()?.get();
  81                 let exchange =
  82                     self.inner.compare_exchange(0, val, Ordering::AcqRel, Ordering::Acquire);
  83                 if let Err(old) = exchange {
  84                     val = old;
  85                 }
  86                 unsafe { NonZeroUsize::new_unchecked(val) }
  87             }
  88         };
  89         Ok(res)
  90     }
  91 }
  92
  93 /// A thread-safe cell which can be written to only once.
  94 #[derive(Default, Debug)]
  95 pub struct OnceBool {
  96     inner: OnceNonZeroUsize,
  97 }
  98
  99 impl OnceBool {
 100     /// Creates a new empty cell.
 101     #[inline]
 102     pub const fn new() -> OnceBool {
 103         OnceBool { inner: OnceNonZeroUsize::new() }
 104     }
 105
 106     /// Gets the underlying value.
 107     #[inline]
 108     pub fn get(&self) -> Option<bool> {
 109         self.inner.get().map(OnceBool::from_usize)
 110     }
 111
 112     /// Sets the contents of this cell to `value`.
 113     ///
 114     /// Returns `Ok(())` if the cell was empty and `Err(())` if it was
 115     /// full.
 116     #[inline]
 117     pub fn set(&self, value: bool) -> Result<(), ()> {
 118         self.inner.set(OnceBool::to_usize(value))
 119     }
 120
 121     /// Gets the contents of the cell, initializing it with `f` if the cell was
 122     /// empty.
 123     ///
 124     /// If several threads concurrently run `get_or_init`, more than one `f` can
 125     /// be called. However, all threads will return the same value, produced by
 126     /// some `f`.
 127     pub fn get_or_init<F>(&self, f: F) -> bool
 128     where
 129         F: FnOnce() -> bool,
 130     {
 131         OnceBool::from_usize(self.inner.get_or_init(|| OnceBool::to_usize(f())))
 132     }
 133
 134     /// Gets the contents of the cell, initializing it with `f` if
 135     /// the cell was empty. If the cell was empty and `f` failed, an
 136     /// error is returned.
 137     ///
 138     /// If several threads concurrently run `get_or_init`, more than one `f` can
 139     /// be called. However, all threads will return the same value, produced by
 140     /// some `f`.
 141     pub fn get_or_try_init<F, E>(&self, f: F) -> Result<bool, E>
 142     where
 143         F: FnOnce() -> Result<bool, E>,
 144     {
 145         self.inner.get_or_try_init(|| f().map(OnceBool::to_usize)).map(OnceBool::from_usize)
 146     }
 147
 148     #[inline]
 149     fn from_usize(value: NonZeroUsize) -> bool {
 150         value.get() == 1
 151     }
 152     #[inline]
 153     fn to_usize(value: bool) -> NonZeroUsize {
 154         unsafe { NonZeroUsize::new_unchecked(if value { 1 } else { 2 }) }
 155     }
 156 }
 157
 158 #[cfg(feature = "alloc")]
 159 pub use self::once_box::OnceBox;
 160
 161 #[cfg(feature = "alloc")]
 162 mod once_box {
 163     use core::{
 164         marker::PhantomData,
 165         ptr,
 166         sync::atomic::{AtomicPtr, Ordering},
 167     };
 168
 169     use alloc::boxed::Box;
 170
 171     /// A thread-safe cell which can be written to only once.
 172     #[derive(Debug)]
 173     pub struct OnceBox<T> {
 174         inner: AtomicPtr<T>,
 175         ghost: PhantomData<Option<Box<T>>>,
 176     }
 177
 178     impl<T> Default for OnceBox<T> {
 179         fn default() -> Self {
 180             Self::new()
 181         }
 182     }
 183
 184     impl<T> Drop for OnceBox<T> {
 185         fn drop(&mut self) {
 186             let ptr = *self.inner.get_mut();
 187             if !ptr.is_null() {
 188                 drop(unsafe { Box::from_raw(ptr) })
 189             }
 190         }
 191     }
 192
 193     impl<T> OnceBox<T> {
 194         /// Creates a new empty cell.
 195         pub const fn new() -> OnceBox<T> {
 196             OnceBox { inner: AtomicPtr::new(ptr::null_mut()), ghost: PhantomData }
 197         }
 198
 199         /// Gets a reference to the underlying value.
 200         pub fn get(&self) -> Option<&T> {
 201             let ptr = self.inner.load(Ordering::Acquire);
 202             if ptr.is_null() {
 203                 return None;
 204             }
 205             Some(unsafe { &*ptr })
 206         }
 207
 208         /// Sets the contents of this cell to `value`.
 209         ///
 210         /// Returns `Ok(())` if the cell was empty and `Err(value)` if it was
 211         /// full.
 212         pub fn set(&self, value: Box<T>) -> Result<(), Box<T>> {
 213             let ptr = Box::into_raw(value);
 214             let exchange = self.inner.compare_exchange(
 215                 ptr::null_mut(),
 216                 ptr,
 217                 Ordering::AcqRel,
 218                 Ordering::Acquire,
 219             );
 220             if let Err(_) = exchange {
 221                 let value = unsafe { Box::from_raw(ptr) };
 222                 return Err(value);
 223             }
 224             Ok(())
 225         }
 226
 227         /// Gets the contents of the cell, initializing it with `f` if the cell was
 228         /// empty.
 229         ///
 230         /// If several threads concurrently run `get_or_init`, more than one `f` can
 231         /// be called. However, all threads will return the same value, produced by
 232         /// some `f`.
 233         pub fn get_or_init<F>(&self, f: F) -> &T
 234         where
 235             F: FnOnce() -> Box<T>,
 236         {
 237             enum Void {}
 238             match self.get_or_try_init(|| Ok::<Box<T>, Void>(f())) {
 239                 Ok(val) => val,
 240                 Err(void) => match void {},
 241             }
 242         }
 243
 244         /// Gets the contents of the cell, initializing it with `f` if
 245         /// the cell was empty. If the cell was empty and `f` failed, an
 246         /// error is returned.
 247         ///
 248         /// If several threads concurrently run `get_or_init`, more than one `f` can
 249         /// be called. However, all threads will return the same value, produced by
 250         /// some `f`.
 251         pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&T, E>
 252         where
 253             F: FnOnce() -> Result<Box<T>, E>,
 254         {
 255             let mut ptr = self.inner.load(Ordering::Acquire);
 256
 257             if ptr.is_null() {
 258                 let val = f()?;
 259                 ptr = Box::into_raw(val);
 260                 let exchange = self.inner.compare_exchange(
 261                     ptr::null_mut(),
 262                     ptr,
 263                     Ordering::AcqRel,
 264                     Ordering::Acquire,
 265                 );
 266                 if let Err(old) = exchange {
 267                     drop(unsafe { Box::from_raw(ptr) });
 268                     ptr = old;
 269                 }
 270             };
 271             Ok(unsafe { &*ptr })
 272         }
 273     }
 274
 275     unsafe impl<T: Sync + Send> Sync for OnceBox<T> {}
 276
 277     /// ```compile_fail
 278     /// struct S(*mut ());
 279     /// unsafe impl Sync for S {}
 280     ///
 281     /// fn share<T: Sync>(_: &T) {}
 282     /// share(&once_cell::race::OnceBox::<S>::new());
 283     /// ```
 284     fn _dummy() {}
 285 }