diff --git a/src/atomic.rs b/src/atomic.rs index e351e3b..a411f17 100644 --- a/src/atomic.rs +++ b/src/atomic.rs @@ -27,12 +27,12 @@ macro_rules! atomic { }; } -pub struct AtomicCell { +pub struct AtomicOption { inner: AtomicCellInner, _phantom: core::marker::PhantomData, } -impl AtomicCell { +impl AtomicOption { pub const fn new() -> Self { Self { inner: AtomicCellInner::none(), @@ -40,10 +40,27 @@ impl AtomicCell { } } + pub fn from_option(value: Option) -> Self { + Self { + inner: AtomicCellInner::from_option(value), + _phantom: core::marker::PhantomData, + } + } + + pub fn into_option(self) -> Option { + self.inner.into_option() + } + pub fn set(&self, value: T) { self.inner.set(value); } + /// Set the value if the cell is empty, returning `Some(value)` if the cell + /// was already occupied. + pub fn try_set(&self, value: T) -> Option { + self.inner.try_set(value) + } + pub fn take(&self) -> Option { self.inner.take() } @@ -77,6 +94,14 @@ impl AtomicCellInner { } } + fn into_option(self) -> Option { + if self.state.load(Ordering::Relaxed) == Self::FULL { + Some(unsafe { ManuallyDrop::into_inner(self.value.into_inner()).assume_init() }) + } else { + None + } + } + fn from_option(value: Option) -> Self { match value { Some(v) => Self { @@ -259,4 +284,92 @@ impl AtomicCellInner { }; None } + + // Try to set the value, returning `value` if the cell already contains a value. + fn try_set(&self, value: T) -> Option { + let this = Self::from_option(Some(value)); + + atomic! { + Self, a, + { + // SAFETY: this block is only executed if `Self` can be transmuted into an atomic type. + // self.state cannot be `LOCKED` here, so we can safely swap the value. + unsafe { + // turn `self` into an atomic pointer + a = &*(self as *const Self as *const _); + + let empty = Self::none(); + // compare-exchange the value with an unset `Self` atomically + if a.compare_exchange(mem::transmute_copy(&empty), mem::transmute_copy(&this), Ordering::Release, Ordering::Relaxed).is_ok() { + None + } else { + this.into_option() + } + } + + }, + { + // Fallback if no atomic type is found. + // we need to lock the cell to swap the value. + + // attempt to lock optimistically + match self.state.compare_exchange_weak( + Self::EMPTY, + Self::LOCKED, + Ordering::Acquire, + Ordering::Relaxed, + ) { + Ok(_) => { + // SAFETY: We are the only thread that can access this cell now. + unsafe { + self.copy_from(&this, Ordering::Relaxed, Ordering::Release); + } + None + } + Err(mut state) => { + let mut spin_wait = SpinWait::new(); + let old = loop { + // if the state is `FULL`, we short-circuit to + // return the provided value. + if state == Self::FULL { + break state; + } + + // if the state is `LOCKED`, we need to wait + if state == Self::LOCKED { + spin_wait.spin(); + continue; + } + + // if the state is not `LOCKED`, we can try locking + // and swapping the value` + if self.state.compare_exchange_weak( + state, + Self::LOCKED, + Ordering::Acquire, + Ordering::Relaxed, + ).is_ok() { + break state; + } else { + // the state changed, we need to check again + state = self.state.load(Ordering::Relaxed); + continue; + } + }; + + if old == Self::EMPTY { + // SAFETY: the cell is locked, so we can safely copy the value + unsafe { + self.copy_from(&this, Ordering::Relaxed, Ordering::Release); + None + } + } else { + this.into_option() + } + } + } + + } + } + } }