RustPython · qingshi163 · Nov 24, 2022 · Nov 13, 2022 · Nov 13, 2022 · Nov 17, 2022
diff --git a/vm/src/object/ext.rs b/vm/src/object/ext.rs
@@ -11,7 +11,7 @@ use crate::{
     convert::{IntoPyException, ToPyObject, ToPyResult, TryFromObject},
     VirtualMachine,
 };
-use std::{borrow::Borrow, fmt, ops::Deref};
+use std::{borrow::Borrow, fmt, marker::PhantomData, ops::Deref, ptr::null_mut};
 
 /* Python objects and references.
 
@@ -214,27 +214,43 @@ impl<T: PyPayload> ToPyObject for PyRefExact<T> {
     }
 }
 
-pub struct PyAtomicRef<T: PyObjectPayload>(PyAtomic<*mut Py<T>>);
+pub struct PyAtomicRef<T> {
+    inner: PyAtomic<*mut u8>,
+    _phantom: PhantomData<T>,
+}
 
 cfg_if::cfg_if! {
     if #[cfg(feature = "threading")] {
         unsafe impl<T: Send + PyObjectPayload> Send for PyAtomicRef<T> {}
         unsafe impl<T: Sync + PyObjectPayload> Sync for PyAtomicRef<T> {}
+        unsafe impl<T: Send + PyObjectPayload> Send for PyAtomicRef<Option<T>> {}
+        unsafe impl<T: Sync + PyObjectPayload> Sync for PyAtomicRef<Option<T>> {}
+        unsafe impl Send for PyAtomicRef<PyObject> {}
+        unsafe impl Sync for PyAtomicRef<PyObject> {}
     }
 }
 
 impl<T: PyObjectPayload> From<PyRef<T>> for PyAtomicRef<T> {
     fn from(pyref: PyRef<T>) -> Self {
         let py = PyRef::leak(pyref);
-        Self(Radium::new(py as *const _ as *mut _))
+        Self {
+            inner: Radium::new(py as *const _ as *mut _),
+            _phantom: Default::default(),
+        }
     }
 }
 
 impl<T: PyObjectPayload> Deref for PyAtomicRef<T> {
     type Target = Py<T>;
 
     fn deref(&self) -> &Self::Target {
-        unsafe { &*self.0.load(Ordering::Relaxed) }
+        unsafe {
+            self.inner
+                .load(Ordering::Relaxed)
+                .cast::<Py<T>>()
+                .as_ref()
+                .unwrap_unchecked()
+        }
     }
 }
 
@@ -244,9 +260,9 @@ impl<T: PyObjectPayload> PyAtomicRef<T> {
     /// until no more reference can be used via PyAtomicRef::deref()
     #[must_use]
     pub unsafe fn swap(&self, pyref: PyRef<T>) -> PyRef<T> {
-        let py = PyRef::leak(pyref);
-        let old = Radium::swap(&self.0, py as *const _ as *mut _, Ordering::AcqRel);
-        PyRef::from_raw(old)
+        let py = PyRef::leak(pyref) as *const Py<T> as *mut _;
+        let old = Radium::swap(&self.inner, py, Ordering::AcqRel);
+        PyRef::from_raw(old.cast())
     }
 
     pub fn swap_to_temporary_refs(&self, pyref: PyRef<T>, vm: &VirtualMachine) {
@@ -257,6 +273,88 @@ impl<T: PyObjectPayload> PyAtomicRef<T> {
     }
 }
 
+impl<T: PyObjectPayload> From<Option<PyRef<T>>> for PyAtomicRef<Option<T>> {
+    fn from(opt_ref: Option<PyRef<T>>) -> Self {
+        let val = opt_ref
+            .map(|x| PyRef::leak(x) as *const Py<T> as *mut _)
+            .unwrap_or(null_mut());
+        Self {
+            inner: Radium::new(val),
+            _phantom: Default::default(),
+        }
+    }
+}
+
+impl<T: PyObjectPayload> PyAtomicRef<Option<T>> {
+    pub fn deref(&self) -> Option<&Py<T>> {
+        unsafe { self.inner.load(Ordering::Relaxed).cast::<Py<T>>().as_ref() }
+    }
+
+    /// # Safety
+    /// The caller is responsible to keep the returned PyRef alive
+    /// until no more reference can be used via PyAtomicRef::deref()
+    #[must_use]
+    pub unsafe fn swap(&self, opt_ref: Option<PyRef<T>>) -> Option<PyRef<T>> {
+        let val = opt_ref
+            .map(|x| PyRef::leak(x) as *const Py<T> as *mut _)
+            .unwrap_or(null_mut());
+        let old = Radium::swap(&self.inner, val, Ordering::AcqRel);
+        unsafe { old.cast::<Py<T>>().as_ref().map(|x| PyRef::from_raw(x)) }
+    }
+
+    pub fn swap_to_temporary_refs(&self, opt_ref: Option<PyRef<T>>, vm: &VirtualMachine) {
+        let Some(old) = (unsafe { self.swap(opt_ref) }) else {
+            return;
+        };
+        if let Some(frame) = vm.current_frame() {
+            frame.temporary_refs.lock().push(old.into());
+        }
+    }
+}
+
+impl From<PyObjectRef> for PyAtomicRef<PyObject> {
+    fn from(obj: PyObjectRef) -> Self {
+        let obj = obj.into_raw();
+        Self {
+            inner: Radium::new(obj as *mut _),
+            _phantom: Default::default(),
+        }
+    }
+}
+
+impl Deref for PyAtomicRef<PyObject> {
+    type Target = PyObject;
+
+    fn deref(&self) -> &Self::Target {
+        unsafe {
+            self.inner
+                .load(Ordering::Relaxed)
+                .cast::<PyObject>()
+                .as_ref()
+                .unwrap_unchecked()
+        }
+    }
+}
+
+impl PyAtomicRef<PyObject> {
+    /// # Safety
+    /// The caller is responsible to keep the returned PyRef alive
+    /// until no more reference can be used via PyAtomicRef::deref()
+    #[must_use]
+    pub unsafe fn swap(&self, obj: PyObjectRef) -> PyObjectRef {
+        let obj = obj.into_raw() as *mut _;
+        let old = Radium::swap(&self.inner, obj, Ordering::AcqRel);
+        PyObjectRef::from_raw(old.cast())
+    }
+
+    pub fn swap_to_temporary_refs(&self, obj: PyObjectRef, vm: &VirtualMachine) {
+        let old = unsafe { self.swap(obj) };
+        if let Some(frame) = vm.current_frame() {
+            frame.temporary_refs.lock().push(old);
+        }
+    }
+}
+
 pub trait AsObject
 where
     Self: Borrow<PyObject>,