split PyCell and PyClassObject concepts (#3917)

* add test for refguard ref counting * split `PyCell` and `PyClassObject` concepts * rework `create_cell` to `create_class_object` * Apply suggestions from code review Co-authored-by: Icxolu <[email protected]> * review: Icxolu feedback --------- Co-authored-by: Icxolu <[email protected]>
PyO3 · Mar 3, 2024 · 2e56f65 · 2e56f65
1 parent 81be11e
commit 2e56f65
Show file tree

Hide file tree

Showing 14 changed files with 301 additions and 294 deletions.
diff --git a/pyo3-macros-backend/src/method.rs b/pyo3-macros-backend/src/method.rs
@@ -645,8 +645,7 @@ impl<'a> FnSpec<'a> {
                         #( #holders )*
                         let result = #call;
                         let initializer: _pyo3::PyClassInitializer::<#cls> = result.convert(py)?;
-                        let cell = initializer.create_cell_from_subtype(py, _slf)?;
-                        ::std::result::Result::Ok(cell as *mut _pyo3::ffi::PyObject)
+                        _pyo3::impl_::pymethods::tp_new_impl(py, initializer, _slf)
                     }
                 }
             }

diff --git a/src/impl_/coroutine.rs b/src/impl_/coroutine.rs
@@ -1,15 +1,15 @@
 use std::{
     future::Future,
-    mem,
     ops::{Deref, DerefMut},
 };
 
 use crate::{
     coroutine::{cancel::ThrowCallback, Coroutine},
     instance::Bound,
+    pycell::impl_::PyClassBorrowChecker,
     pyclass::boolean_struct::False,
     types::{PyAnyMethods, PyString},
-    IntoPy, Py, PyAny, PyCell, PyClass, PyErr, PyObject, PyResult, Python,
+    IntoPy, Py, PyAny, PyClass, PyErr, PyObject, PyResult, Python,
 };
 
 pub fn new_coroutine<F, T, E>(
@@ -32,17 +32,16 @@ where
 }
 
 fn get_ptr<T: PyClass>(obj: &Py<T>) -> *mut T {
-    // SAFETY: Py<T> can be casted as *const PyCell<T>
-    unsafe { &*(obj.as_ptr() as *const PyCell<T>) }.get_ptr()
+    obj.get_class_object().get_ptr()
 }
 
 pub struct RefGuard<T: PyClass>(Py<T>);
 
 impl<T: PyClass> RefGuard<T> {
     pub fn new(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
-        let owned = obj.downcast::<T>()?;
-        mem::forget(owned.try_borrow()?);
-        Ok(RefGuard(owned.clone().unbind()))
+        let bound = obj.downcast::<T>()?;
+        bound.get_class_object().borrow_checker().try_borrow()?;
+        Ok(RefGuard(bound.clone().unbind()))
     }
 }
 
@@ -57,9 +56,11 @@ impl<T: PyClass> Deref for RefGuard<T> {
 impl<T: PyClass> Drop for RefGuard<T> {
     fn drop(&mut self) {
         Python::with_gil(|gil| {
-            #[allow(deprecated)]
-            let self_ref = self.0.bind(gil);
-            self_ref.release_ref()
+            self.0
+                .bind(gil)
+                .get_class_object()
+                .borrow_checker()
+                .release_borrow()
         })
     }
 }
@@ -68,9 +69,9 @@ pub struct RefMutGuard<T: PyClass<Frozen = False>>(Py<T>);
 
 impl<T: PyClass<Frozen = False>> RefMutGuard<T> {
     pub fn new(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
-        let owned = obj.downcast::<T>()?;
-        mem::forget(owned.try_borrow_mut()?);
-        Ok(RefMutGuard(owned.clone().unbind()))
+        let bound = obj.downcast::<T>()?;
+        bound.get_class_object().borrow_checker().try_borrow_mut()?;
+        Ok(RefMutGuard(bound.clone().unbind()))
     }
 }
 
@@ -92,9 +93,11 @@ impl<T: PyClass<Frozen = False>> DerefMut for RefMutGuard<T> {
 impl<T: PyClass<Frozen = False>> Drop for RefMutGuard<T> {
     fn drop(&mut self) {
         Python::with_gil(|gil| {
-            #[allow(deprecated)]
-            let self_ref = self.0.bind(gil);
-            self_ref.release_mut()
+            self.0
+                .bind(gil)
+                .get_class_object()
+                .borrow_checker()
+                .release_borrow_mut()
         })
     }
 }
diff --git a/src/impl_/pycell.rs b/src/impl_/pycell.rs
@@ -1,2 +1,4 @@
 //! Externally-accessible implementation of pycell
-pub use crate::pycell::impl_::{GetBorrowChecker, PyClassMutability};
+pub use crate::pycell::impl_::{
+    GetBorrowChecker, PyClassMutability, PyClassObject, PyClassObjectBase, PyClassObjectLayout,
+};
diff --git a/src/impl_/pyclass.rs b/src/impl_/pyclass.rs
@@ -2,14 +2,13 @@ use crate::{
     exceptions::{PyAttributeError, PyNotImplementedError, PyRuntimeError, PyValueError},
     ffi,
     impl_::freelist::FreeList,
-    impl_::pycell::{GetBorrowChecker, PyClassMutability},
+    impl_::pycell::{GetBorrowChecker, PyClassMutability, PyClassObjectLayout},
     internal_tricks::extract_c_string,
-    pycell::PyCellLayout,
     pyclass_init::PyObjectInit,
     types::any::PyAnyMethods,
     types::PyBool,
-    Borrowed, Py, PyAny, PyCell, PyClass, PyErr, PyMethodDefType, PyNativeType, PyResult,
-    PyTypeInfo, Python,
+    Borrowed, Py, PyAny, PyClass, PyErr, PyMethodDefType, PyNativeType, PyResult, PyTypeInfo,
+    Python,
 };
 use std::{
     borrow::Cow,
@@ -26,13 +25,13 @@ pub use lazy_type_object::LazyTypeObject;
 /// Gets the offset of the dictionary from the start of the object in bytes.
 #[inline]
 pub fn dict_offset<T: PyClass>() -> ffi::Py_ssize_t {
-    PyCell::<T>::dict_offset()
+    PyClassObject::<T>::dict_offset()
 }
 
 /// Gets the offset of the weakref list from the start of the object in bytes.
 #[inline]
 pub fn weaklist_offset<T: PyClass>() -> ffi::Py_ssize_t {
-    PyCell::<T>::weaklist_offset()
+    PyClassObject::<T>::weaklist_offset()
 }
 
 /// Represents the `__dict__` field for `#[pyclass]`.
@@ -883,6 +882,8 @@ macro_rules! generate_pyclass_richcompare_slot {
 }
 pub use generate_pyclass_richcompare_slot;
 
+use super::pycell::PyClassObject;
+
 /// Implements a freelist.
 ///
 /// Do not implement this trait manually. Instead, use `#[pyclass(freelist = N)]`
@@ -1095,7 +1096,7 @@ impl<T> PyClassThreadChecker<T> for ThreadCheckerImpl {
 
 /// Trait denoting that this class is suitable to be used as a base type for PyClass.
 pub trait PyClassBaseType: Sized {
-    type LayoutAsBase: PyCellLayout<Self>;
+    type LayoutAsBase: PyClassObjectLayout<Self>;
     type BaseNativeType;
     type Initializer: PyObjectInit<Self>;
     type PyClassMutability: PyClassMutability;
@@ -1105,15 +1106,15 @@ pub trait PyClassBaseType: Sized {
 ///
 /// In the future this will be extended to immutable PyClasses too.
 impl<T: PyClass> PyClassBaseType for T {
-    type LayoutAsBase = crate::pycell::PyCell<T>;
+    type LayoutAsBase = crate::impl_::pycell::PyClassObject<T>;
     type BaseNativeType = T::BaseNativeType;
     type Initializer = crate::pyclass_init::PyClassInitializer<Self>;
     type PyClassMutability = T::PyClassMutability;
 }
 
 /// Implementation of tp_dealloc for pyclasses without gc
 pub(crate) unsafe extern "C" fn tp_dealloc<T: PyClass>(obj: *mut ffi::PyObject) {
-    crate::impl_::trampoline::dealloc(obj, PyCell::<T>::tp_dealloc)
+    crate::impl_::trampoline::dealloc(obj, PyClassObject::<T>::tp_dealloc)
 }
 
 /// Implementation of tp_dealloc for pyclasses with gc
@@ -1122,7 +1123,7 @@ pub(crate) unsafe extern "C" fn tp_dealloc_with_gc<T: PyClass>(obj: *mut ffi::Py
     {
         ffi::PyObject_GC_UnTrack(obj.cast());
     }
-    crate::impl_::trampoline::dealloc(obj, PyCell::<T>::tp_dealloc)
+    crate::impl_::trampoline::dealloc(obj, PyClassObject::<T>::tp_dealloc)
 }
 
 pub(crate) unsafe extern "C" fn get_sequence_item_from_mapping(

diff --git a/src/impl_/pymethods.rs b/src/impl_/pymethods.rs
@@ -7,8 +7,8 @@ use crate::pycell::{PyBorrowError, PyBorrowMutError};
 use crate::pyclass::boolean_struct::False;
 use crate::types::{any::PyAnyMethods, PyModule, PyType};
 use crate::{
-    ffi, Borrowed, Bound, DowncastError, Py, PyAny, PyCell, PyClass, PyErr, PyObject, PyRef,
-    PyRefMut, PyResult, PyTraverseError, PyTypeCheck, PyVisit, Python,
+    ffi, Borrowed, Bound, DowncastError, Py, PyAny, PyCell, PyClass, PyClassInitializer, PyErr,
+    PyObject, PyRef, PyRefMut, PyResult, PyTraverseError, PyTypeCheck, PyVisit, Python,
 };
 use std::borrow::Cow;
 use std::ffi::CStr;
@@ -569,3 +569,13 @@ impl<'py, T> std::ops::Deref for BoundRef<'_, 'py, T> {
         self.0
     }
 }
+
+pub unsafe fn tp_new_impl<T: PyClass>(
+    py: Python<'_>,
+    initializer: PyClassInitializer<T>,
+    target_type: *mut ffi::PyTypeObject,
+) -> PyResult<*mut ffi::PyObject> {
+    initializer
+        .create_class_object_of_type(py, target_type)
+        .map(Bound::into_ptr)
+}
diff --git a/src/instance.rs b/src/instance.rs
@@ -1,5 +1,5 @@
 use crate::err::{self, PyDowncastError, PyErr, PyResult};
-use crate::ffi_ptr_ext::FfiPtrExt;
+use crate::impl_::pycell::PyClassObject;
 use crate::pycell::{PyBorrowError, PyBorrowMutError, PyCell};
 use crate::pyclass::boolean_struct::{False, True};
 use crate::type_object::HasPyGilRef;
@@ -92,14 +92,7 @@ where
         py: Python<'py>,
         value: impl Into<PyClassInitializer<T>>,
     ) -> PyResult<Bound<'py, T>> {
-        let initializer = value.into();
-        let obj = initializer.create_cell(py)?;
-        let ob = unsafe {
-            obj.cast::<ffi::PyObject>()
-                .assume_owned(py)
-                .downcast_into_unchecked()
-        };
-        Ok(ob)
+        value.into().create_class_object(py)
     }
 }
 
@@ -332,19 +325,11 @@ where
     where
         T: PyClass<Frozen = True> + Sync,
     {
-        let cell = self.get_cell();
-        // SAFETY: The class itself is frozen and `Sync` and we do not access anything but `cell.contents.value`.
-        unsafe { &*cell.get_ptr() }
+        self.1.get()
     }
 
-    pub(crate) fn get_cell(&'py self) -> &'py PyCell<T> {
-        let cell = self.as_ptr().cast::<PyCell<T>>();
-        // SAFETY: Bound<T> is known to contain an object which is laid out in memory as a
-        // PyCell<T>.
-        //
-        // Strictly speaking for now `&'py PyCell<T>` is part of the "GIL Ref" API, so this
-        // could use some further refactoring later to avoid going through this reference.
-        unsafe { &*cell }
+    pub(crate) fn get_class_object(&self) -> &PyClassObject<T> {
+        self.1.get_class_object()
     }
 }
 
@@ -887,10 +872,7 @@ where
     /// # }
     /// ```
     pub fn new(py: Python<'_>, value: impl Into<PyClassInitializer<T>>) -> PyResult<Py<T>> {
-        let initializer = value.into();
-        let obj = initializer.create_cell(py)?;
-        let ob = unsafe { Py::from_owned_ptr(py, obj as _) };
-        Ok(ob)
+        Bound::new(py, value).map(Bound::unbind)
     }
 }
 
@@ -1194,12 +1176,16 @@ where
     where
         T: PyClass<Frozen = True> + Sync,
     {
-        let any = self.as_ptr() as *const PyAny;
-        // SAFETY: The class itself is frozen and `Sync` and we do not access anything but `cell.contents.value`.
-        unsafe {
-            let cell: &PyCell<T> = PyNativeType::unchecked_downcast(&*any);
-            &*cell.get_ptr()
-        }
+        // Safety: The class itself is frozen and `Sync`
+        unsafe { &*self.get_class_object().get_ptr() }
+    }
+
+    /// Get a view on the underlying `PyClass` contents.
+    pub(crate) fn get_class_object(&self) -> &PyClassObject<T> {
+        let class_object = self.as_ptr().cast::<PyClassObject<T>>();
+        // Safety: Bound<T: PyClass> is known to contain an object which is laid out in memory as a
+        // PyClassObject<T>.
+        unsafe { &*class_object }
     }
 }