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 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] new() -> OnceNonZeroUsize23 pub const fn new() -> OnceNonZeroUsize { 24 OnceNonZeroUsize { inner: AtomicUsize::new(0) } 25 } 26 27 /// Gets the underlying value. 28 #[inline] get(&self) -> Option<NonZeroUsize>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] set(&self, value: NonZeroUsize) -> Result<(), ()>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`. get_or_init<F>(&self, f: F) -> NonZeroUsize where F: FnOnce() -> NonZeroUsize,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`. get_or_try_init<F, E>(&self, f: F) -> Result<NonZeroUsize, E> where F: FnOnce() -> Result<NonZeroUsize, E>,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] new() -> OnceBool102 pub const fn new() -> OnceBool { 103 OnceBool { inner: OnceNonZeroUsize::new() } 104 } 105 106 /// Gets the underlying value. 107 #[inline] get(&self) -> Option<bool>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] set(&self, value: bool) -> Result<(), ()>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`. get_or_init<F>(&self, f: F) -> bool where F: FnOnce() -> bool,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`. get_or_try_init<F, E>(&self, f: F) -> Result<bool, E> where F: FnOnce() -> Result<bool, E>,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] from_usize(value: NonZeroUsize) -> bool149 fn from_usize(value: NonZeroUsize) -> bool { 150 value.get() == 1 151 } 152 #[inline] to_usize(value: bool) -> NonZeroUsize153 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> { default() -> Self179 fn default() -> Self { 180 Self::new() 181 } 182 } 183 184 impl<T> Drop for OnceBox<T> { drop(&mut self)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. new() -> OnceBox<T>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. get(&self) -> Option<&T>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. set(&self, value: Box<T>) -> Result<(), Box<T>>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`. get_or_init<F>(&self, f: F) -> &T where F: FnOnce() -> Box<T>,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`. get_or_try_init<F, E>(&self, f: F) -> Result<&T, E> where F: FnOnce() -> Result<Box<T>, E>,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 /// ``` _dummy()284 fn _dummy() {} 285 } 286