AtomicOption try_set, into/from_option

src/atomic.rs

@@ -27,12 +27,12 @@ macro_rules! atomic {
     };
 }
 
-pub struct AtomicCell<T> {
+pub struct AtomicOption<T> {
     inner: AtomicCellInner<T>,
     _phantom: core::marker::PhantomData<T>,
 }
 
-impl<T> AtomicCell<T> {
+impl<T> AtomicOption<T> {
     pub const fn new() -> Self {
         Self {
             inner: AtomicCellInner::none(),
@@ -40,10 +40,27 @@ impl<T> AtomicCell<T> {
         }
     }
 
+    pub fn from_option(value: Option<T>) -> Self {
+        Self {
+            inner: AtomicCellInner::from_option(value),
+            _phantom: core::marker::PhantomData,
+        }
+    }
+
+    pub fn into_option(self) -> Option<T> {
+        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<T> {
+        self.inner.try_set(value)
+    }
+
     pub fn take(&self) -> Option<T> {
         self.inner.take()
     }
@@ -77,6 +94,14 @@ impl<T> AtomicCellInner<T> {
         }
     }
 
+    fn into_option(self) -> Option<T> {
+        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<T>) -> Self {
         match value {
             Some(v) => Self {
@@ -259,4 +284,92 @@ impl<T> AtomicCellInner<T> {
         };
         None
     }
+
+    // Try to set the value, returning `value` if the cell already contains a value.
+    fn try_set(&self, value: T) -> Option<T> {
+        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()
+                        }
+                    }
+                }
+
+            }
+        }
+    }
 }