mindustry logic execution, map- and schematic- parsing and rendering
Add a cell which can be initialized atomically
| -rw-r--r-- | src/main.rs | 1 | ||||
| -rw-r--r-- | src/utils/mod.rs | 2 | ||||
| -rw-r--r-- | src/utils/once_cell.rs | 172 |
3 files changed, 175 insertions, 0 deletions
diff --git a/src/main.rs b/src/main.rs index 17f2d2c..24edbc3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,6 +10,7 @@ pub mod modifier; pub mod registry; pub mod team; pub mod unit; +pub mod utils; fn main() { diff --git a/src/utils/mod.rs b/src/utils/mod.rs new file mode 100644 index 0000000..396b1e2 --- /dev/null +++ b/src/utils/mod.rs @@ -0,0 +1,2 @@ +pub mod once_cell; +pub use once_cell::OnceCell; diff --git a/src/utils/once_cell.rs b/src/utils/once_cell.rs new file mode 100644 index 0000000..a5e72f7 --- /dev/null +++ b/src/utils/once_cell.rs @@ -0,0 +1,172 @@ +use std::cell::UnsafeCell; +use std::mem::MaybeUninit; +use std::ptr; +use std::sync::atomic::{AtomicU8, Ordering}; + +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, +} + +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> 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> 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> {} + +unsafe impl<T: Sync> Sync for OnceCell<T> {} |