WIP: AtomicPy

Author: Icxolu

src/atomic.rs

@@ -0,0 +1,197 @@
+use std::{
+    marker::PhantomData,
+    sync::atomic::{AtomicPtr, Ordering},
+};
+
+use crate::{ffi, Borrowed, Bound, Py, Python};
+
+/// Variant of [`Py<T>`] that can be atomically swapped
+#[repr(transparent)]
+pub struct AtomicPy<T> {
+    inner: AtomicPtr<ffi::PyObject>,
+    _marker: PhantomData<T>,
+}
+
+// SAFETY: same as `Py<T>`
+unsafe impl<T> Send for AtomicPy<T> {}
+unsafe impl<T> Sync for AtomicPy<T> {}
+
+impl<T> AtomicPy<T> {
+    /// Create an [`AtomicPy<T>`] holding the [`Py<T>`]
+    #[inline]
+    pub fn from_py(obj: Py<T>) -> Self {
+        Self {
+            inner: AtomicPtr::new(obj.into_ptr()),
+            _marker: PhantomData,
+        }
+    }
+
+    /// Create an [`AtomicPy<T>`] holding the [`Bound<T>`]
+    #[inline]
+    pub fn from_bound(obj: Bound<'_, T>) -> Self {
+        Self::from_py(obj.unbind())
+    }
+
+    /// Consumes the [`AtomicPy<T>`] and turns it back into a [`Py<T>`]
+    #[inline]
+    pub fn into_inner(self) -> Py<T> {
+        let ptr = *std::mem::ManuallyDrop::new(self).inner.get_mut();
+        unsafe { Py::from_owned_ptr_unchecked(ptr) }
+    }
+
+    /// Loads the object stored in the [`AtomicPy`]
+    #[inline]
+    pub fn load<'py>(&self, py: Python<'py>, order: Ordering) -> Bound<'py, T> {
+        let ptr = self.inner.load(order);
+        unsafe {
+            Borrowed::from_ptr_unchecked(py, ptr)
+                .to_owned()
+                .cast_into_unchecked()
+        }
+    }
+
+    /// Stores `obj` in the [`AtomicPy`], dropping its current value
+    ///
+    /// Note: This uses [`swap`](Self::swap) under the hood.
+    #[inline]
+    pub fn store<'py>(&self, obj: Bound<'py, T>, order: Ordering) {
+        let _ = self.swap(obj, order);
+    }
+
+    /// Stores `obj` in the [`AtomicPy`], returning the previous value.
+    ///
+    /// See [`swap`](AtomicPy::swap)
+    #[inline]
+    pub fn swap_unbound(&self, obj: Py<T>, order: Ordering) -> Py<T> {
+        let ptr = self.inner.swap(obj.into_ptr(), order);
+        // SAFETY: `ptr` is a non-null, owned pointer to `ffi::PyObject` of type `T`
+        unsafe { Py::from_owned_ptr_unchecked(ptr) }
+    }
+
+    /// Stores `obj` in the [`AtomicPy`], returning the previous value.
+    ///
+    /// `swap` takes an [Ordering] argument which describes the memory ordering of this operation.
+    /// All ordering modes are possible. Note that using [Acquire](Ordering::Acquire) makes the
+    /// store part of this operation [Relaxed](Ordering::Relaxed), and using
+    /// [Release](Ordering::Release) makes the load part [Relaxed](Ordering::Relaxed).
+    #[inline]
+    pub fn swap<'py>(&self, obj: Bound<'py, T>, order: Ordering) -> Bound<'py, T> {
+        let py = obj.py();
+        self.swap_unbound(obj.unbind(), order).into_bound(py)
+    }
+
+    /// Stores `new` into the [`AtomicPy`] if its current value is `current` (Python `current is
+    /// self`)
+    ///
+    /// The return value is a [`Result`] indicating whether `new` was written and containing the
+    /// previous value.
+    ///
+    /// `compare_exchange` takes two [Ordering] arguments to describe the memory ordering of this
+    /// operation. `success` describes the required ordering for the read-modify-write operation
+    /// that takes place if the comparison with `current` succeeds. `failure` describes the required
+    /// ordering for the load operation that takes place when the comparison fails. Using
+    /// [Acquire](Ordering::Acquire) as success ordering makes the store part of this operation
+    /// [Relaxed](Ordering::Relaxed), and using [Release](Ordering::Release) makes the successful
+    /// load [Relaxed](Ordering::Relaxed). The failure ordering can only be
+    /// [SeqCst](Ordering::SeqCst), [Acquire](Ordering::Acquire) or [Relaxed](Ordering::Relaxed).
+    #[inline]
+    pub fn compare_exchange<'py>(
+        &self,
+        current: &Bound<'py, T>,
+        new: &Bound<'py, T>,
+        success: Ordering,
+        failure: Ordering,
+    ) -> Result<Bound<'py, T>, Bound<'py, T>> {
+        let py = current.py();
+        match self
+            .inner
+            .compare_exchange(current.as_ptr(), new.as_ptr(), success, failure)
+        {
+            // stored `new`, `ptr` is owned
+            Ok(ptr) => unsafe {
+                let _ = std::mem::ManuallyDrop::new(new.clone()); // only incref if `new` was successfully stored
+
+                // SAFETY: `ptr` is a non-null, owned pointer to `ffi::PyObject` of type `T`
+                Ok(Bound::from_owned_ptr_unchecked(py, ptr).cast_into_unchecked())
+            },
+            // did not store `new`, `ptr` is borrowed
+            Err(ptr) => unsafe {
+                // SAFETY: `ptr` is a non-null, borrowed pointer to `ffi::PyObject` of type `T`
+                Err(Borrowed::from_ptr_unchecked(py, ptr)
+                    .to_owned()
+                    .cast_into_unchecked())
+            },
+        }
+    }
+
+    /// Stores `new` into the [`AtomicPy`] if its current value is `current` (Python `current is
+    /// self`)
+    ///
+    /// Unlike [AtomicPy::compare_exchange], this function is allowed to spuriously fail even when
+    /// the comparison succeeds, which can result in more efficient code on some platforms. The
+    /// return value is a [`Result`] indicating whether `new` was written and containing the
+    /// previous value.
+    ///
+    /// `compare_exchange_weak` takes two [Ordering] arguments to describe the memory ordering of
+    /// this operation. `success` describes the required ordering for the read-modify-write
+    /// operation that takes place if the comparison with `current` succeeds. `failure` describes
+    /// the required ordering for the load operation that takes place when the comparison fails.
+    /// Using [Acquire](Ordering::Acquire) as success ordering makes the store part of this
+    /// operation [Relaxed](Ordering::Relaxed), and using [Release](Ordering::Release) makes the
+    /// successful load [Relaxed](Ordering::Relaxed). The failure ordering can only be
+    /// [SeqCst](Ordering::SeqCst), [Acquire](Ordering::Acquire) or [Relaxed](Ordering::Relaxed).
+    #[inline]
+    pub fn compare_exchange_weak<'py>(
+        &self,
+        current: &Bound<'py, T>,
+        new: &Bound<'py, T>,
+        success: Ordering,
+        failure: Ordering,
+    ) -> Result<Bound<'py, T>, Bound<'py, T>> {
+        let py = current.py();
+        match self
+            .inner
+            .compare_exchange_weak(current.as_ptr(), new.as_ptr(), success, failure)
+        {
+            // stored `new`, `ptr` is owned `current`
+            Ok(ptr) => unsafe {
+                let _ = std::mem::ManuallyDrop::new(new.clone()); // only incref if `new` was successfully stored
+
+                // SAFETY: `ptr` is a non-null, owned pointer to `ffi::PyObject` of type `T`
+                Ok(Bound::from_owned_ptr_unchecked(py, ptr).cast_into_unchecked())
+            },
+            // did not store `new`, `ptr` is borrowed
+            Err(ptr) => unsafe {
+                // SAFETY: `ptr` is a non-null, borrowed pointer to `ffi::PyObject` of type `T`
+                Err(Borrowed::from_ptr_unchecked(py, ptr)
+                    .to_owned()
+                    .cast_into_unchecked())
+            },
+        }
+    }
+}
+
+impl<T> From<Py<T>> for AtomicPy<T> {
+    /// Create an [`AtomicPy<T>`] holding the [`Py<T>`]
+    fn from(obj: Py<T>) -> Self {
+        Self::from_py(obj)
+    }
+}
+
+impl<T> From<Bound<'_, T>> for AtomicPy<T> {
+    /// Create an [`AtomicPy<T>`] holding the [`Bound<T>`]
+    fn from(obj: Bound<'_, T>) -> Self {
+        Self::from_bound(obj)
+    }
+}
+
+impl<T> Drop for AtomicPy<T> {
+    /// Drop the inner [`Py<T>`]
+    fn drop(&mut self) {
+        // SAFETY: `inner` is a non-null, owned pointer to `ffi::PyObject` of type `T`
+        unsafe { Py::<T>::from_owned_ptr_unchecked(*self.inner.get_mut()) };
+    }
+}
+
+#[cfg(test)]
+mod tests {}

src/lib.rs

@@ -422,6 +422,7 @@ pub mod impl_;
 mod internal_tricks;
 mod internal;
 
+mod atomic;
 pub mod buffer;
 pub mod call;
 pub mod conversion;
@@ -449,6 +450,7 @@ mod version;
 
 #[allow(unused_imports)] // with no features enabled this module has no public exports
 pub use crate::conversions::*;
+pub use atomic::AtomicPy;
 
 #[cfg(feature = "macros")]
 pub use pyo3_macros::{