HeapTypeExt::type_data

youknowone · Dec 11, 2025 · Nov 30, 2025 · Dec 11, 2025 · Dec 11, 2025 · 81fa1f4b70e43f302850089caee341b8a7410729
commit 81fa1f4b70e43f302850089caee341b8a7410729
diff --git a/crates/vm/src/builtins/type.rs b/crates/vm/src/builtins/type.rs
@@ -26,12 +26,13 @@ use crate::{
     protocol::{PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods},
     types::{
         AsNumber, Callable, Constructor, GetAttr, PyTypeFlags, PyTypeSlots, Representable, SetAttr,
+        TypeDataRef, TypeDataRefMut, TypeDataSlot,
     },
 };
 use indexmap::{IndexMap, map::Entry};
 use itertools::Itertools;
 use num_traits::ToPrimitive;
-use std::{borrow::Borrow, collections::HashSet, ops::Deref, pin::Pin, ptr::NonNull};
+use std::{any::Any, borrow::Borrow, collections::HashSet, ops::Deref, pin::Pin, ptr::NonNull};
 
 #[pyclass(module = false, name = "type", traverse = "manual")]
 pub struct PyType {
@@ -65,6 +66,7 @@ pub struct HeapTypeExt {
     pub slots: Option<PyRef<PyTuple<PyStrRef>>>,
     pub sequence_methods: PySequenceMethods,
     pub mapping_methods: PyMappingMethods,
+    pub type_data: PyRwLock<Option<TypeDataSlot>>,
 }
 
 pub struct PointerSlot<T>(NonNull<T>);
@@ -203,6 +205,7 @@ impl PyType {
             slots: None,
             sequence_methods: PySequenceMethods::default(),
             mapping_methods: PyMappingMethods::default(),
+            type_data: PyRwLock::new(None),
         };
         let base = bases[0].clone();
 
@@ -563,6 +566,50 @@ impl PyType {
             |ext| PyRwLockReadGuard::map(ext.name.read(), |name| name.as_str()).into(),
         )
     }
+
+    // Type Data Slot API - CPython's PyObject_GetTypeData equivalent
+
+    /// Initialize type data for this type. Can only be called once.
+    /// Returns an error if the type is not a heap type or if data is already initialized.
+    pub fn init_type_data<T: Any + Send + Sync + 'static>(&self, data: T) -> Result<(), String> {
+        let ext = self
+            .heaptype_ext
+            .as_ref()
+            .ok_or_else(|| "Cannot set type data on non-heap types".to_string())?;
+
+        let mut type_data = ext.type_data.write();
+        if type_data.is_some() {
+            return Err("Type data already initialized".to_string());
+        }
+        *type_data = Some(TypeDataSlot::new(data));
+        Ok(())
+    }
+
+    /// Get a read guard to the type data.
+    /// Returns None if the type is not a heap type, has no data, or the data type doesn't match.
+    pub fn get_type_data<T: Any + 'static>(&self) -> Option<TypeDataRef<'_, T>> {
+        self.heaptype_ext
+            .as_ref()
+            .and_then(|ext| TypeDataRef::try_new(ext.type_data.read()))
+    }
+
+    /// Get a write guard to the type data.
+    /// Returns None if the type is not a heap type, has no data, or the data type doesn't match.
+    pub fn get_type_data_mut<T: Any + 'static>(&self) -> Option<TypeDataRefMut<'_, T>> {
+        self.heaptype_ext
+            .as_ref()
+            .and_then(|ext| TypeDataRefMut::try_new(ext.type_data.write()))
+    }
+
+    /// Check if this type has type data of the given type.
+    pub fn has_type_data<T: Any + 'static>(&self) -> bool {
+        self.heaptype_ext.as_ref().is_some_and(|ext| {
+            ext.type_data
+                .read()
+                .as_ref()
+                .is_some_and(|slot| slot.get::<T>().is_some())
+        })
+    }
 }
 
 impl Py<PyType> {
@@ -1167,6 +1214,7 @@ impl Constructor for PyType {
                 slots: heaptype_slots.clone(),
                 sequence_methods: PySequenceMethods::default(),
                 mapping_methods: PyMappingMethods::default(),
+                type_data: PyRwLock::new(None),
             };
             (slots, heaptype_ext)
         };

diff --git a/crates/vm/src/types/slot.rs b/crates/vm/src/types/slot.rs
@@ -1,3 +1,6 @@
+use crate::common::lock::{
+    PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyRwLockReadGuard, PyRwLockWriteGuard,
+};
 use crate::{
     AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
     builtins::{PyInt, PyStr, PyStrInterned, PyStrRef, PyType, PyTypeRef, type_::PointerSlot},
@@ -16,7 +19,96 @@ use crate::{
 use crossbeam_utils::atomic::AtomicCell;
 use malachite_bigint::BigInt;
 use num_traits::{Signed, ToPrimitive};
-use std::{borrow::Borrow, cmp::Ordering, ops::Deref};
+use std::{any::Any, any::TypeId, borrow::Borrow, cmp::Ordering, ops::Deref};
+
+/// Type-erased storage for extension module data attached to heap types.
+pub struct TypeDataSlot {
+    // PyObject_GetTypeData
+    type_id: TypeId,
+    data: Box<dyn Any + Send + Sync>,
+}
+
+impl TypeDataSlot {
+    /// Create a new type data slot with the given data.
+    pub fn new<T: Any + Send + Sync + 'static>(data: T) -> Self {
+        Self {
+            type_id: TypeId::of::<T>(),
+            data: Box::new(data),
+        }
+    }
+
+    /// Get a reference to the data if the type matches.
+    pub fn get<T: Any + 'static>(&self) -> Option<&T> {
+        if self.type_id == TypeId::of::<T>() {
+            self.data.downcast_ref()
+        } else {
+            None
+        }
+    }
+
+    /// Get a mutable reference to the data if the type matches.
+    pub fn get_mut<T: Any + 'static>(&mut self) -> Option<&mut T> {
+        if self.type_id == TypeId::of::<T>() {
+            self.data.downcast_mut()
+        } else {
+            None
+        }
+    }
+}
+
+/// Read guard for type data access, using mapped guard for zero-cost deref.
+pub struct TypeDataRef<'a, T: 'static> {
+    guard: PyMappedRwLockReadGuard<'a, T>,
+}
+
+impl<'a, T: Any + 'static> TypeDataRef<'a, T> {
+    /// Try to create a TypeDataRef from a read guard.
+    /// Returns None if the slot is empty or contains a different type.
+    pub fn try_new(guard: PyRwLockReadGuard<'a, Option<TypeDataSlot>>) -> Option<Self> {
+        PyRwLockReadGuard::try_map(guard, |opt| opt.as_ref().and_then(|slot| slot.get::<T>()))
+            .ok()
+            .map(|guard| Self { guard })
+    }
+}
+
+impl<T: Any + 'static> std::ops::Deref for TypeDataRef<'_, T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.guard
+    }
+}
+
+/// Write guard for type data access, using mapped guard for zero-cost deref.
+pub struct TypeDataRefMut<'a, T: 'static> {
+    guard: PyMappedRwLockWriteGuard<'a, T>,
+}
+
+impl<'a, T: Any + 'static> TypeDataRefMut<'a, T> {
+    /// Try to create a TypeDataRefMut from a write guard.
+    /// Returns None if the slot is empty or contains a different type.
+    pub fn try_new(guard: PyRwLockWriteGuard<'a, Option<TypeDataSlot>>) -> Option<Self> {
+        PyRwLockWriteGuard::try_map(guard, |opt| {
+            opt.as_mut().and_then(|slot| slot.get_mut::<T>())
+        })
+        .ok()
+        .map(|guard| Self { guard })
+    }
+}
+
+impl<T: Any + 'static> std::ops::Deref for TypeDataRefMut<'_, T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.guard
+    }
+}
+
+impl<T: Any + 'static> std::ops::DerefMut for TypeDataRefMut<'_, T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.guard
+    }
+}
 
 #[macro_export]
 macro_rules! atomic_func {