introduce fast path to LazyTypeObject::get_or_init

newsfragments/5324.changed.md

@@ -0,0 +1 @@
+Add fast-path to `PyTypeInfo::type_object` for `#[pyclass]` types.

src/impl_/pyclass/lazy_type_object.rs

@@ -11,12 +11,11 @@ use crate::types::PyTypeMethods;
 use crate::{
     exceptions::PyRuntimeError,
     ffi,
-    impl_::pyclass::MaybeRuntimePyMethodDef,
-    impl_::pymethods::PyMethodDefType,
+    impl_::{pyclass::MaybeRuntimePyMethodDef, pymethods::PyMethodDefType},
     pyclass::{create_type_object, PyClassTypeObject},
     sync::GILOnceCell,
     types::PyType,
-    Bound, PyClass, PyErr, PyObject, PyResult, Python,
+    Bound, Py, PyClass, PyErr, PyObject, PyResult, Python,
 };
 
 use std::sync::Mutex;
@@ -33,7 +32,7 @@ struct LazyTypeObjectInner {
     // Threads which have begun initialization of the `tp_dict`. Used for
     // reentrant initialization detection.
     initializing_threads: Mutex<Vec<ThreadId>>,
-    tp_dict_filled: GILOnceCell<()>,
+    fully_initialized_type: GILOnceCell<Py<PyType>>,
 }
 
 impl<T> LazyTypeObject<T> {
@@ -44,7 +43,7 @@ impl<T> LazyTypeObject<T> {
             LazyTypeObjectInner {
                 value: GILOnceCell::new(),
                 initializing_threads: Mutex::new(Vec::new()),
-                tp_dict_filled: GILOnceCell::new(),
+                fully_initialized_type: GILOnceCell::new(),
             },
             PhantomData,
         )
@@ -53,15 +52,37 @@ impl<T> LazyTypeObject<T> {
 
 impl<T: PyClass> LazyTypeObject<T> {
     /// Gets the type object contained by this `LazyTypeObject`, initializing it if needed.
+    #[inline]
     pub fn get_or_init<'py>(&self, py: Python<'py>) -> &Bound<'py, PyType> {
-        self.get_or_try_init(py).unwrap_or_else(|err| {
+        if let Some(type_object) = self.0.fully_initialized_type.get(py) {
+            // Fast path
+            return type_object.bind(py);
+        }
+
+        self.init(py)
+    }
+
+    /// Fallible version of the above.
+    #[inline]
+    pub(crate) fn get_or_try_init<'py>(&self, py: Python<'py>) -> PyResult<&Bound<'py, PyType>> {
+        if let Some(type_object) = self.0.fully_initialized_type.get(py) {
+            // Fast path
+            return Ok(type_object.bind(py));
+        }
+
+        self.try_init(py)
+    }
+
+    #[cold]
+    fn init<'py>(&self, py: Python<'py>) -> &Bound<'py, PyType> {
+        self.try_init(py).unwrap_or_else(|err| {
             err.print(py);
             panic!("failed to create type object for {}", T::NAME)
         })
     }
 
-    /// Fallible version of the above.
-    pub(crate) fn get_or_try_init<'py>(&self, py: Python<'py>) -> PyResult<&Bound<'py, PyType>> {
+    #[cold]
+    fn try_init<'py>(&self, py: Python<'py>) -> PyResult<&Bound<'py, PyType>> {
         self.0
             .get_or_try_init(py, create_type_object::<T>, T::NAME, T::items_iter())
     }
@@ -117,7 +138,7 @@ impl LazyTypeObjectInner {
         // `tp_dict`, it can still request the type object through `get_or_init`,
         // but the `tp_dict` may appear empty of course.
 
-        if self.tp_dict_filled.get(py).is_some() {
+        if self.fully_initialized_type.get(py).is_some() {
             // `tp_dict` is already filled: ok.
             return Ok(());
         }
@@ -185,8 +206,8 @@ impl LazyTypeObjectInner {
 
         // Now we hold the GIL and we can assume it won't be released until we
         // return from the function.
-        let result = self.tp_dict_filled.get_or_try_init(py, move || {
-            let result = initialize_tp_dict(py, type_object.as_ptr(), items);
+        let result = self.fully_initialized_type.get_or_try_init(py, move || {
+            initialize_tp_dict(py, type_object.as_ptr(), items)?;
             #[cfg(Py_3_14)]
             if is_immutable_type {
                 // freeze immutable types after __dict__ is initialized
@@ -217,13 +238,13 @@ impl LazyTypeObjectInner {
                 self.initializing_threads.lock().unwrap()
             };
             threads.clear();
-            result
+            Ok(type_object.clone().unbind())
         });
 
         if let Err(err) = result {
             return Err(wrap_in_runtime_error(
                 py,
-                err.clone_ref(py),
+                err,
                 format!("An error occurred while initializing `{name}.__dict__`"),
             ));
         }