Skip to content

Commit d34b2cf

Browse files
youknowoneyouknowone
authored
TypeData (#6403)
* HeapTypeExt::type_data * Apply TypeDataSlot to ctypes
1 parent bfd873a commit d34b2cf

File tree

9 files changed

+348
-180
lines changed

9 files changed

+348
-180
lines changed

‎crates/vm/src/builtins/type.rs‎

Lines changed: 49 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,12 +26,13 @@ use crate::{
2626
protocol::{PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods},
2727
types::{
2828
AsNumber, Callable, Constructor, GetAttr, PyTypeFlags, PyTypeSlots, Representable, SetAttr,
29+
TypeDataRef, TypeDataRefMut, TypeDataSlot,
2930
},
3031
};
3132
use indexmap::{IndexMap, map::Entry};
3233
use itertools::Itertools;
3334
use num_traits::ToPrimitive;
34-
use std::{borrow::Borrow, collections::HashSet, ops::Deref, pin::Pin, ptr::NonNull};
35+
use std::{any::Any, borrow::Borrow, collections::HashSet, ops::Deref, pin::Pin, ptr::NonNull};
3536

3637
#[pyclass(module = false, name = "type", traverse = "manual")]
3738
pub struct PyType {
@@ -65,6 +66,7 @@ pub struct HeapTypeExt {
6566
pub slots: Option<PyRef<PyTuple<PyStrRef>>>,
6667
pub sequence_methods: PySequenceMethods,
6768
pub mapping_methods: PyMappingMethods,
69+
pub type_data: PyRwLock<Option<TypeDataSlot>>,
6870
}
6971

7072
pub struct PointerSlot<T>(NonNull<T>);
@@ -203,6 +205,7 @@ impl PyType {
203205
slots: None,
204206
sequence_methods: PySequenceMethods::default(),
205207
mapping_methods: PyMappingMethods::default(),
208+
type_data: PyRwLock::new(None),
206209
};
207210
let base = bases[0].clone();
208211

@@ -563,6 +566,50 @@ impl PyType {
563566
|ext| PyRwLockReadGuard::map(ext.name.read(), |name| name.as_str()).into(),
564567
)
565568
}
569+
570+
// Type Data Slot API - CPython's PyObject_GetTypeData equivalent
571+
572+
/// Initialize type data for this type. Can only be called once.
573+
/// Returns an error if the type is not a heap type or if data is already initialized.
574+
pub fn init_type_data<T: Any + Send + Sync + 'static>(&self, data: T) -> Result<(), String> {
575+
let ext = self
576+
.heaptype_ext
577+
.as_ref()
578+
.ok_or_else(|| "Cannot set type data on non-heap types".to_string())?;
579+
580+
let mut type_data = ext.type_data.write();
581+
if type_data.is_some() {
582+
return Err("Type data already initialized".to_string());
583+
}
584+
*type_data = Some(TypeDataSlot::new(data));
585+
Ok(())
586+
}
587+
588+
/// Get a read guard to the type data.
589+
/// Returns None if the type is not a heap type, has no data, or the data type doesn't match.
590+
pub fn get_type_data<T: Any + 'static>(&self) -> Option<TypeDataRef<'_, T>> {
591+
self.heaptype_ext
592+
.as_ref()
593+
.and_then(|ext| TypeDataRef::try_new(ext.type_data.read()))
594+
}
595+
596+
/// Get a write guard to the type data.
597+
/// Returns None if the type is not a heap type, has no data, or the data type doesn't match.
598+
pub fn get_type_data_mut<T: Any + 'static>(&self) -> Option<TypeDataRefMut<'_, T>> {
599+
self.heaptype_ext
600+
.as_ref()
601+
.and_then(|ext| TypeDataRefMut::try_new(ext.type_data.write()))
602+
}
603+
604+
/// Check if this type has type data of the given type.
605+
pub fn has_type_data<T: Any + 'static>(&self) -> bool {
606+
self.heaptype_ext.as_ref().is_some_and(|ext| {
607+
ext.type_data
608+
.read()
609+
.as_ref()
610+
.is_some_and(|slot| slot.get::<T>().is_some())
611+
})
612+
}
566613
}
567614

568615
impl Py<PyType> {
@@ -1167,6 +1214,7 @@ impl Constructor for PyType {
11671214
slots: heaptype_slots.clone(),
11681215
sequence_methods: PySequenceMethods::default(),
11691216
mapping_methods: PyMappingMethods::default(),
1217+
type_data: PyRwLock::new(None),
11701218
};
11711219
(slots, heaptype_ext)
11721220
};

‎crates/vm/src/stdlib/ctypes.rs‎

Lines changed: 29 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -389,39 +389,35 @@ pub(crate) mod _ctypes {
389389
/// Get the size of a ctypes type or instance
390390
#[pyfunction(name = "sizeof")]
391391
pub fn size_of(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<usize> {
392-
use super::array::{PyCArray, PyCArrayType};
393392
use super::pointer::PyCPointer;
394393
use super::structure::{PyCStructType, PyCStructure};
395-
use super::union::{PyCUnion, PyCUnionType};
394+
use super::union::PyCUnionType;
395+
use super::util::StgInfo;
396+
use crate::builtins::PyType;
396397

397-
// 1. Instances with stg_info
398-
if obj.fast_isinstance(PyCArray::static_type()) {
399-
// Get stg_info from the type
400-
if let Some(type_obj) = obj.class().as_object().downcast_ref::<PyCArrayType>() {
401-
return Ok(type_obj.stg_info.size);
402-
}
398+
// 1. Check TypeDataSlot on class (for instances)
399+
if let Some(stg_info) = obj.class().get_type_data::<StgInfo>() {
400+
return Ok(stg_info.size);
403401
}
402+
403+
// 2. Check TypeDataSlot on type itself (for type objects)
404+
if let Some(type_obj) = obj.downcast_ref::<PyType>()
405+
&& let Some(stg_info) = type_obj.get_type_data::<StgInfo>()
406+
{
407+
return Ok(stg_info.size);
408+
}
409+
410+
// 3. Instances with cdata buffer
404411
if let Some(structure) = obj.downcast_ref::<PyCStructure>() {
405412
return Ok(structure.cdata.read().size());
406413
}
407-
if obj.fast_isinstance(PyCUnion::static_type()) {
408-
// Get stg_info from the type
409-
if let Some(type_obj) = obj.class().as_object().downcast_ref::<PyCUnionType>() {
410-
return Ok(type_obj.stg_info.size);
411-
}
412-
}
413414
if let Some(simple) = obj.downcast_ref::<PyCSimple>() {
414415
return Ok(simple.cdata.read().size());
415416
}
416417
if obj.fast_isinstance(PyCPointer::static_type()) {
417418
return Ok(std::mem::size_of::<usize>());
418419
}
419420

420-
// 2. Types (metatypes with stg_info)
421-
if let Some(array_type) = obj.downcast_ref::<PyCArrayType>() {
422-
return Ok(array_type.stg_info.size);
423-
}
424-
425421
// 3. Type objects
426422
if let Ok(type_ref) = obj.clone().downcast::<crate::builtins::PyType>() {
427423
// Structure types - check if metaclass is or inherits from PyCStructType
@@ -659,33 +655,37 @@ pub(crate) mod _ctypes {
659655

660656
#[pyfunction]
661657
fn alignment(tp: Either<PyTypeRef, PyObjectRef>, vm: &VirtualMachine) -> PyResult<usize> {
662-
use super::array::{PyCArray, PyCArrayType};
663658
use super::base::PyCSimpleType;
664659
use super::pointer::PyCPointer;
665660
use super::structure::PyCStructure;
666661
use super::union::PyCUnion;
662+
use super::util::StgInfo;
663+
use crate::builtins::PyType;
667664

668665
let obj = match &tp {
669666
Either::A(t) => t.as_object(),
670667
Either::B(o) => o.as_ref(),
671668
};
672669

673-
// Try to get alignment from stg_info directly (for instances)
674-
if let Some(array_type) = obj.downcast_ref::<PyCArrayType>() {
675-
return Ok(array_type.stg_info.align);
670+
// 1. Check TypeDataSlot on class (for instances)
671+
if let Some(stg_info) = obj.class().get_type_data::<StgInfo>() {
672+
return Ok(stg_info.align);
676673
}
674+
675+
// 2. Check TypeDataSlot on type itself (for type objects)
676+
if let Some(type_obj) = obj.downcast_ref::<PyType>()
677+
&& let Some(stg_info) = type_obj.get_type_data::<StgInfo>()
678+
{
679+
return Ok(stg_info.align);
680+
}
681+
682+
// 3. Fallback for simple types without TypeDataSlot
677683
if obj.fast_isinstance(PyCSimple::static_type()) {
678684
// Get stg_info from the type by reading _type_ attribute
679685
let cls = obj.class().to_owned();
680686
let stg_info = PyCSimpleType::get_stg_info(&cls, vm);
681687
return Ok(stg_info.align);
682688
}
683-
if obj.fast_isinstance(PyCArray::static_type()) {
684-
// Get stg_info from the type
685-
if let Some(type_obj) = obj.class().as_object().downcast_ref::<PyCArrayType>() {
686-
return Ok(type_obj.stg_info.align);
687-
}
688-
}
689689
if obj.fast_isinstance(PyCStructure::static_type()) {
690690
// Calculate alignment from _fields_
691691
let cls = obj.class();

0 commit comments

Comments
 (0)