atomic_option.rs

use std::{
    cell::Cell,
    sync::{Arc, RwLock}, thread,
};

#[derive(Debug)]
pub struct AtomicOption<T> {
    value: Arc<RwLock<Cell<*mut T>>>,
}

unsafe impl<T> Send for AtomicOption<T> {}
unsafe impl<T> Sync for AtomicOption<T> {}

impl<T> AtomicOption<T> {
    pub fn new(value: Option<Arc<T>>) -> Self {
        let result = AtomicOption::empty();
        result.store(value);
        result
    }

    pub fn empty() -> Self {
        Self {
            value: Arc::new(RwLock::new(Cell::new(std::ptr::null_mut()))),
        }
    }

    /// Gets the underlying value.
    pub fn load(&self) -> Option<Arc<T>> {
        let guard = self.value.read().unwrap();
        let ptr = guard.get();
        if ptr.is_null() {
            None
        } else {
            // SAFETY: The current arc will not drop and its reference count will increment.
            unsafe {
                // do not call drop as to prevent a decrement
                let arc = std::mem::ManuallyDrop::new(Arc::<T>::from_raw(ptr));
                // clone to increment
                Some(std::mem::ManuallyDrop::into_inner(arc.clone()))
            }
        }
    }

    /// Sets the underlying value.
    pub fn store(&self, value: Option<Arc<T>>) {
        let ptr = match value {
            Some(value) => Arc::into_raw(value) as *mut T,
            None => std::ptr::null_mut(),
        };

        let old_ptr = {
            let guard = self.value.write().unwrap();
            let old_ptr = guard.get();
            guard.set(ptr);
            old_ptr
        };
        if !old_ptr.is_null() {
            // SAFETY: Decrements the old arc's reference count by dropping it.
            unsafe { Arc::from_raw(old_ptr) };
        }
    }
}

impl<T> Drop for AtomicOption<T> {
    fn drop(&mut self) {
        self.store(None);
    }
}