diff --git a/vm/src/object/ext.rs b/vm/src/object/ext.rs index 9a0e9acbd7..d9683b35d9 100644 --- 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,19 +214,29 @@ impl ToPyObject for PyRefExact { } } -pub struct PyAtomicRef(PyAtomic<*mut Py>); +pub struct PyAtomicRef { + inner: PyAtomic<*mut u8>, + _phantom: PhantomData, +} cfg_if::cfg_if! { if #[cfg(feature = "threading")] { unsafe impl Send for PyAtomicRef {} unsafe impl Sync for PyAtomicRef {} + unsafe impl Send for PyAtomicRef> {} + unsafe impl Sync for PyAtomicRef> {} + unsafe impl Send for PyAtomicRef {} + unsafe impl Sync for PyAtomicRef {} } } impl From> for PyAtomicRef { fn from(pyref: PyRef) -> 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(), + } } } @@ -234,7 +244,13 @@ impl Deref for PyAtomicRef { type Target = Py; fn deref(&self) -> &Self::Target { - unsafe { &*self.0.load(Ordering::Relaxed) } + unsafe { + self.inner + .load(Ordering::Relaxed) + .cast::>() + .as_ref() + .unwrap_unchecked() + } } } @@ -244,9 +260,9 @@ impl PyAtomicRef { /// until no more reference can be used via PyAtomicRef::deref() #[must_use] pub unsafe fn swap(&self, pyref: PyRef) -> PyRef { - 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 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, vm: &VirtualMachine) { @@ -257,6 +273,88 @@ impl PyAtomicRef { } } +impl From>> for PyAtomicRef> { + fn from(opt_ref: Option>) -> Self { + let val = opt_ref + .map(|x| PyRef::leak(x) as *const Py as *mut _) + .unwrap_or(null_mut()); + Self { + inner: Radium::new(val), + _phantom: Default::default(), + } + } +} + +impl PyAtomicRef> { + pub fn deref(&self) -> Option<&Py> { + unsafe { self.inner.load(Ordering::Relaxed).cast::>().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>) -> Option> { + let val = opt_ref + .map(|x| PyRef::leak(x) as *const Py as *mut _) + .unwrap_or(null_mut()); + let old = Radium::swap(&self.inner, val, Ordering::AcqRel); + unsafe { old.cast::>().as_ref().map(|x| PyRef::from_raw(x)) } + } + + pub fn swap_to_temporary_refs(&self, opt_ref: Option>, 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 for PyAtomicRef { + fn from(obj: PyObjectRef) -> Self { + let obj = obj.into_raw(); + Self { + inner: Radium::new(obj as *mut _), + _phantom: Default::default(), + } + } +} + +impl Deref for PyAtomicRef { + type Target = PyObject; + + fn deref(&self) -> &Self::Target { + unsafe { + self.inner + .load(Ordering::Relaxed) + .cast::() + .as_ref() + .unwrap_unchecked() + } + } +} + +impl PyAtomicRef { + /// # 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,