mindustry logic execution, map- and schematic- parsing and rendering
Diffstat (limited to 'src/utils/once_cell.rs')
| -rw-r--r-- | src/utils/once_cell.rs | 292 |
1 files changed, 144 insertions, 148 deletions
diff --git a/src/utils/once_cell.rs b/src/utils/once_cell.rs index a5e72f7..4bfd3eb 100644 --- a/src/utils/once_cell.rs +++ b/src/utils/once_cell.rs @@ -7,164 +7,160 @@ const STATE_INIT: u8 = 0; const STATE_LOCKED: u8 = STATE_INIT + 1; const STATE_READY: u8 = STATE_LOCKED + 1; -pub struct OnceCell<T> -{ - value: UnsafeCell<MaybeUninit<T>>, - state: AtomicU8, +pub struct OnceCell<T> { + value: UnsafeCell<MaybeUninit<T>>, + state: AtomicU8, } -impl<T> OnceCell<T> -{ - pub const fn new() -> Self - { - Self{value: UnsafeCell::new(MaybeUninit::uninit()), state: AtomicU8::new(STATE_INIT)} - } - - pub const fn new_init(value: T) -> Self - { - Self{value: UnsafeCell::new(MaybeUninit::new(value)), state: AtomicU8::new(STATE_READY)} - } - - pub fn get(&self) -> Option<&T> - { - if self.state.load(Ordering::Acquire) == STATE_READY - { - // SAFETY: won't be overwritten for the lifetime of this reference - Some(unsafe{(&*self.value.get()).assume_init_ref()}) - } - else {None} - } - - pub fn get_or_wait(&self) -> Option<&T> - { - loop - { - match self.state.load(Ordering::Acquire) - { - STATE_INIT => return None, - STATE_LOCKED => (), // continue - STATE_READY => return Some(unsafe{(&*self.value.get()).assume_init_ref()}), - x => unreachable!("invalid state {x}"), - } - } - } - - pub fn get_or_init<F: Fn() -> T>(&self, init: F) -> &T - { - loop - { - match self.state.compare_exchange(STATE_INIT, STATE_LOCKED, Ordering::AcqRel, Ordering::Acquire) - { - Ok(..) => - { - let value = init(); - let written = &*unsafe{&mut *self.value.get()}.write(value); - self.state.store(STATE_READY, Ordering::Release); - return written; - }, - Err(STATE_READY) => return unsafe{(&*self.value.get()).assume_init_ref()}, - Err(..) => (), // locked or spurious failure - } - } - } - - pub fn set(&self, value: T) -> Result<&T, T> - { - // don't set state to STATE_READY on success because we have to release afterward anyway - match self.state.compare_exchange(STATE_INIT, STATE_LOCKED, Ordering::AcqRel, Ordering::Acquire) - { - Ok(..) => - { - // SAFETY: unique because only one thread can lock the atomic state - let written = &*unsafe{&mut *self.value.get()}.write(value); - self.state.store(STATE_READY, Ordering::Release); - Ok(written) - }, - // SAFETY: guaranteed to be initialized & protected by acquire ordering - Err(STATE_READY) => return Ok(unsafe{(&*self.value.get()).assume_init_ref()}), - Err(..) => Err(value), // locked or spurious failure - } - } - - pub fn set_mut(&mut self, mut value: T) -> Result<Option<T>, T> - { - // don't set state to STATE_READY on success because we have to release afterward anyway - match self.state.compare_exchange(STATE_INIT, STATE_LOCKED, Ordering::AcqRel, Ordering::Acquire) - { - Ok(..) => - { - self.value.get_mut().write(value); - self.state.store(STATE_READY, Ordering::Release); - Ok(None) - }, - Err(STATE_READY) => - { - // SAFETY: guaranteed to be initialized & protected by acquire ordering - std::mem::swap(unsafe{self.value.get_mut().assume_init_mut()}, &mut value); - // ensure changes are visible to others acquiring the atomic state - self.state.store(STATE_READY, Ordering::Release); - // we've swapped the previous value into this variable - Ok(Some(value)) - }, - Err(..) => Err(value), // locked or spurious failure - } - } - - pub fn into_inner(mut self) -> Option<T> - { - // must be atomic so potential writes during the drop see a valid state - let inner = match self.state.load(Ordering::Acquire) - { - STATE_INIT => None, - STATE_LOCKED => unreachable!("consumed cell during initialization"), - // SAFETY: initialized & we'll forget about it afterwards - STATE_READY => Some(unsafe{self.value.get_mut().assume_init_read()}), - x => unreachable!("invalid state {x}"), - }; - // SAFETY: just in case AtomicU8 has a drop handler - unsafe{ptr::drop_in_place(&mut self.state as *mut _);} - std::mem::forget(self); - inner - } +impl<T> OnceCell<T> { + pub const fn new() -> Self { + Self { + value: UnsafeCell::new(MaybeUninit::uninit()), + state: AtomicU8::new(STATE_INIT), + } + } + + pub const fn new_init(value: T) -> Self { + Self { + value: UnsafeCell::new(MaybeUninit::new(value)), + state: AtomicU8::new(STATE_READY), + } + } + + pub fn get(&self) -> Option<&T> { + if self.state.load(Ordering::Acquire) == STATE_READY { + // SAFETY: won't be overwritten for the lifetime of this reference + Some(unsafe { (&*self.value.get()).assume_init_ref() }) + } else { + None + } + } + + pub fn get_or_wait(&self) -> Option<&T> { + loop { + match self.state.load(Ordering::Acquire) { + STATE_INIT => return None, + STATE_LOCKED => (), // continue + STATE_READY => return Some(unsafe { (&*self.value.get()).assume_init_ref() }), + x => unreachable!("invalid state {x}"), + } + } + } + + pub fn get_or_init<F: Fn() -> T>(&self, init: F) -> &T { + loop { + match self.state.compare_exchange( + STATE_INIT, + STATE_LOCKED, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(..) => { + let value = init(); + let written = &*unsafe { &mut *self.value.get() }.write(value); + self.state.store(STATE_READY, Ordering::Release); + return written; + } + Err(STATE_READY) => return unsafe { (&*self.value.get()).assume_init_ref() }, + Err(..) => (), // locked or spurious failure + } + } + } + + pub fn set(&self, value: T) -> Result<&T, T> { + // don't set state to STATE_READY on success because we have to release afterward anyway + match self.state.compare_exchange( + STATE_INIT, + STATE_LOCKED, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(..) => { + // SAFETY: unique because only one thread can lock the atomic state + let written = &*unsafe { &mut *self.value.get() }.write(value); + self.state.store(STATE_READY, Ordering::Release); + Ok(written) + } + // SAFETY: guaranteed to be initialized & protected by acquire ordering + Err(STATE_READY) => return Ok(unsafe { (&*self.value.get()).assume_init_ref() }), + Err(..) => Err(value), // locked or spurious failure + } + } + + pub fn set_mut(&mut self, mut value: T) -> Result<Option<T>, T> { + // don't set state to STATE_READY on success because we have to release afterward anyway + match self.state.compare_exchange( + STATE_INIT, + STATE_LOCKED, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(..) => { + self.value.get_mut().write(value); + self.state.store(STATE_READY, Ordering::Release); + Ok(None) + } + Err(STATE_READY) => { + // SAFETY: guaranteed to be initialized & protected by acquire ordering + std::mem::swap( + unsafe { self.value.get_mut().assume_init_mut() }, + &mut value, + ); + // ensure changes are visible to others acquiring the atomic state + self.state.store(STATE_READY, Ordering::Release); + // we've swapped the previous value into this variable + Ok(Some(value)) + } + Err(..) => Err(value), // locked or spurious failure + } + } + + pub fn into_inner(mut self) -> Option<T> { + // must be atomic so potential writes during the drop see a valid state + let inner = match self.state.load(Ordering::Acquire) { + STATE_INIT => None, + STATE_LOCKED => unreachable!("consumed cell during initialization"), + // SAFETY: initialized & we'll forget about it afterwards + STATE_READY => Some(unsafe { self.value.get_mut().assume_init_read() }), + x => unreachable!("invalid state {x}"), + }; + // SAFETY: just in case AtomicU8 has a drop handler + unsafe { + ptr::drop_in_place(&mut self.state as *mut _); + } + std::mem::forget(self); + inner + } } -impl<T> Default for OnceCell<T> -{ - fn default() -> Self - { - OnceCell::new() - } +impl<T> Default for OnceCell<T> { + fn default() -> Self { + OnceCell::new() + } } -impl<T> From<T> for OnceCell<T> -{ - fn from(value: T) -> Self - { - OnceCell::new_init(value) - } +impl<T> From<T> for OnceCell<T> { + fn from(value: T) -> Self { + OnceCell::new_init(value) + } } -impl<T> From<OnceCell<T>> for Option<T> -{ - fn from(value: OnceCell<T>) -> Self - { - value.into_inner() - } +impl<T> From<OnceCell<T>> for Option<T> { + fn from(value: OnceCell<T>) -> Self { + value.into_inner() + } } -impl<T> Drop for OnceCell<T> -{ - fn drop(&mut self) - { - match *self.state.get_mut() - { - STATE_INIT => (), - STATE_LOCKED => unreachable!("dropped cell during initialization"), - // MaybeUninit requires us to manually drop the value - STATE_READY => unsafe{self.value.get_mut().assume_init_drop()}, - x => unreachable!("invalid state {x}"), - } - } +impl<T> Drop for OnceCell<T> { + fn drop(&mut self) { + match *self.state.get_mut() { + STATE_INIT => (), + STATE_LOCKED => unreachable!("dropped cell during initialization"), + // MaybeUninit requires us to manually drop the value + STATE_READY => unsafe { self.value.get_mut().assume_init_drop() }, + x => unreachable!("invalid state {x}"), + } + } } unsafe impl<T: Send> Send for OnceCell<T> {} |