🌐 AI搜索 & 代理 主页
Skip to content

Commit 0ab4805

Browse files
youknowoneyouknowone
committed
Apply TypeDataSlot to ctypes
1 parent 3077e89 commit 0ab4805

File tree

7 files changed

+206
-178
lines changed

7 files changed

+206
-178
lines changed

‎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();

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

Lines changed: 87 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
use crate::atomic_func;
22
use crate::builtins::{PyBytes, PyInt};
3-
use crate::convert::ToPyObject;
3+
use crate::class::StaticType;
44
use crate::function::FuncArgs;
55
use crate::protocol::{
66
BufferDescriptor, BufferMethods, PyBuffer, PyNumberMethods, PySequenceMethods,
77
};
88
use crate::stdlib::ctypes::base::CDataObject;
99
use crate::stdlib::ctypes::util::StgInfo;
10-
use crate::types::{AsBuffer, AsNumber, AsSequence, Callable};
10+
use crate::types::{AsBuffer, AsNumber, AsSequence};
1111
use crate::{AsObject, Py, PyObjectRef, PyPayload};
1212
use crate::{
1313
PyResult, VirtualMachine,
@@ -20,56 +20,49 @@ use rustpython_common::lock::PyRwLock;
2020
use rustpython_vm::stdlib::ctypes::_ctypes::get_size;
2121
use rustpython_vm::stdlib::ctypes::base::PyCData;
2222

23+
/// PyCArrayType - metatype for Array types
24+
/// CPython stores array info (type, length) in StgInfo via type_data
2325
#[pyclass(name = "PyCArrayType", base = PyType, module = "_ctypes")]
24-
#[derive(PyPayload)]
25-
pub struct PyCArrayType {
26-
pub(super) stg_info: StgInfo,
27-
pub(super) typ: PyRwLock<PyObjectRef>,
28-
pub(super) length: AtomicCell<usize>,
29-
pub(super) element_size: AtomicCell<usize>,
30-
}
31-
32-
impl std::fmt::Debug for PyCArrayType {
33-
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34-
f.debug_struct("PyCArrayType")
35-
.field("typ", &self.typ)
36-
.field("length", &self.length)
37-
.finish()
38-
}
39-
}
40-
41-
impl Callable for PyCArrayType {
42-
type Args = FuncArgs;
43-
fn call(zelf: &Py<Self>, args: Self::Args, vm: &VirtualMachine) -> PyResult {
44-
// Create an instance of the array
45-
let element_type = zelf.typ.read().clone();
46-
let length = zelf.length.load();
47-
let element_size = zelf.element_size.load();
48-
let total_size = element_size * length;
49-
let mut buffer = vec![0u8; total_size];
50-
51-
// Initialize from positional arguments
52-
for (i, value) in args.args.iter().enumerate() {
53-
if i >= length {
54-
break;
55-
}
56-
let offset = i * element_size;
57-
if let Ok(int_val) = value.try_int(vm) {
58-
let bytes = PyCArray::int_to_bytes(int_val.as_bigint(), element_size);
59-
if offset + element_size <= buffer.len() {
60-
buffer[offset..offset + element_size].copy_from_slice(&bytes);
61-
}
62-
}
26+
#[derive(Debug, Default, PyPayload)]
27+
pub struct PyCArrayType {}
28+
29+
/// Create a new Array type with StgInfo stored in type_data (CPython style)
30+
pub fn create_array_type_with_stginfo(stg_info: StgInfo, vm: &VirtualMachine) -> PyResult {
31+
// Get PyCArrayType as metaclass
32+
let metaclass = PyCArrayType::static_type().to_owned();
33+
34+
// Create a unique name for the array type
35+
let type_name = format!("Array_{}", stg_info.length);
36+
37+
// Create args for type(): (name, bases, dict)
38+
let name = vm.ctx.new_str(type_name);
39+
let bases = vm
40+
.ctx
41+
.new_tuple(vec![PyCArray::static_type().to_owned().into()]);
42+
let dict = vm.ctx.new_dict();
43+
44+
let args = FuncArgs::new(
45+
vec![name.into(), bases.into(), dict.into()],
46+
crate::function::KwArgs::default(),
47+
);
48+
49+
// Create the new type using PyType::slot_new with PyCArrayType as metaclass
50+
let new_type = crate::builtins::type_::PyType::slot_new(metaclass, args, vm)?;
51+
52+
// Set StgInfo in type_data
53+
let type_ref: PyTypeRef = new_type
54+
.clone()
55+
.downcast()
56+
.map_err(|_| vm.new_type_error("Failed to create array type".to_owned()))?;
57+
58+
if type_ref.init_type_data(stg_info.clone()).is_err() {
59+
// Type data already initialized - update it
60+
if let Some(mut existing) = type_ref.get_type_data_mut::<StgInfo>() {
61+
*existing = stg_info;
6362
}
64-
65-
Ok(PyCArray {
66-
typ: PyRwLock::new(element_type),
67-
length: AtomicCell::new(length),
68-
element_size: AtomicCell::new(element_size),
69-
cdata: PyRwLock::new(CDataObject::from_bytes(buffer, None)),
70-
}
71-
.into_pyobject(vm))
7263
}
64+
65+
Ok(new_type)
7366
}
7467

7568
impl Constructor for PyCArrayType {
@@ -80,54 +73,62 @@ impl Constructor for PyCArrayType {
8073
}
8174
}
8275

83-
#[pyclass(flags(IMMUTABLETYPE), with(Callable, Constructor, AsNumber))]
76+
#[pyclass(flags(IMMUTABLETYPE), with(Constructor, AsNumber))]
8477
impl PyCArrayType {
8578
#[pygetset(name = "_type_")]
86-
fn typ(&self) -> PyObjectRef {
87-
self.typ.read().clone()
79+
fn typ(zelf: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
80+
zelf.downcast_ref::<PyType>()
81+
.and_then(|t| t.get_type_data::<StgInfo>())
82+
.and_then(|stg| stg.element_type.clone())
83+
.unwrap_or_else(|| vm.ctx.none())
8884
}
8985

9086
#[pygetset(name = "_length_")]
91-
fn length(&self) -> usize {
92-
self.length.load()
87+
fn length(zelf: PyObjectRef) -> usize {
88+
zelf.downcast_ref::<PyType>()
89+
.and_then(|t| t.get_type_data::<StgInfo>())
90+
.map(|stg| stg.length)
91+
.unwrap_or(0)
9392
}
9493

9594
#[pymethod]
96-
fn __mul__(zelf: &Py<Self>, n: isize, vm: &VirtualMachine) -> PyResult {
95+
fn __mul__(zelf: PyObjectRef, n: isize, vm: &VirtualMachine) -> PyResult {
9796
if n < 0 {
9897
return Err(vm.new_value_error(format!("Array length must be >= 0, not {n}")));
9998
}
100-
// Create a nested array type: (inner_type * inner_length) * n
101-
// The new array has n elements, each element is the current array type
102-
// e.g., (c_int * 5) * 3 = Array of 3 elements, each is (c_int * 5)
103-
let inner_length = zelf.length.load();
104-
let inner_element_size = zelf.element_size.load();
99+
100+
// Get inner array info from TypeDataSlot
101+
let type_ref = zelf.downcast_ref::<PyType>().unwrap();
102+
let (_inner_length, inner_size) = type_ref
103+
.get_type_data::<StgInfo>()
104+
.map(|stg| (stg.length, stg.size))
105+
.unwrap_or((0, 0));
105106

106107
// The element type of the new array is the current array type itself
107-
let current_array_type: PyObjectRef = zelf.as_object().to_owned();
108+
let current_array_type: PyObjectRef = zelf.clone();
108109

109110
// Element size is the total size of the inner array
110-
let new_element_size = inner_length * inner_element_size;
111+
let new_element_size = inner_size;
111112
let total_size = new_element_size * (n as usize);
112-
let stg_info = StgInfo::new(total_size, inner_element_size);
113113

114-
Ok(PyCArrayType {
115-
stg_info,
116-
typ: PyRwLock::new(current_array_type),
117-
length: AtomicCell::new(n as usize),
118-
element_size: AtomicCell::new(new_element_size),
119-
}
120-
.to_pyobject(vm))
114+
let stg_info = StgInfo::new_array(
115+
total_size,
116+
new_element_size,
117+
n as usize,
118+
current_array_type,
119+
new_element_size,
120+
);
121+
122+
create_array_type_with_stginfo(stg_info, vm)
121123
}
122124

123125
#[pyclassmethod]
124126
fn in_dll(
125-
zelf: &Py<Self>,
127+
zelf: PyObjectRef,
126128
dll: PyObjectRef,
127129
name: crate::builtins::PyStrRef,
128130
vm: &VirtualMachine,
129131
) -> PyResult {
130-
use crate::stdlib::ctypes::_ctypes::size_of;
131132
use libloading::Symbol;
132133

133134
// Get the library handle from dll object
@@ -168,10 +169,18 @@ impl PyCArrayType {
168169
return Err(vm.new_attribute_error("Library is closed".to_owned()));
169170
};
170171

171-
// Get size from the array type
172-
let element_type = zelf.typ.read().clone();
173-
let length = zelf.length.load();
174-
let element_size = size_of(element_type.clone(), vm)?;
172+
// Get size from the array type via TypeDataSlot
173+
let type_ref = zelf.downcast_ref::<PyType>().unwrap();
174+
let (element_type, length, element_size) = type_ref
175+
.get_type_data::<StgInfo>()
176+
.map(|stg| {
177+
(
178+
stg.element_type.clone().unwrap_or_else(|| vm.ctx.none()),
179+
stg.length,
180+
stg.element_size,
181+
)
182+
})
183+
.unwrap_or_else(|| (vm.ctx.none(), 0, 0));
175184
let total_size = element_size * length;
176185

177186
// Read data from symbol address
@@ -206,15 +215,13 @@ impl AsNumber for PyCArrayType {
206215
fn as_number() -> &'static PyNumberMethods {
207216
static AS_NUMBER: PyNumberMethods = PyNumberMethods {
208217
multiply: Some(|a, b, vm| {
209-
let zelf = a
210-
.downcast_ref::<PyCArrayType>()
211-
.ok_or_else(|| vm.new_type_error("expected PyCArrayType".to_owned()))?;
218+
// a is a type object whose metaclass is PyCArrayType (e.g., Array_5)
212219
let n = b
213220
.try_index(vm)?
214221
.as_bigint()
215222
.to_isize()
216223
.ok_or_else(|| vm.new_overflow_error("array size too large".to_owned()))?;
217-
PyCArrayType::__mul__(zelf, n, vm)
224+
PyCArrayType::__mul__(a.to_owned(), n, vm)
218225
}),
219226
..PyNumberMethods::NOT_IMPLEMENTED
220227
};

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

Lines changed: 11 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,6 @@
11
use super::_ctypes::bytes_to_pyobject;
2-
use super::array::PyCArrayType;
32
use super::util::StgInfo;
43
use crate::builtins::{PyBytes, PyFloat, PyInt, PyNone, PyStr, PyStrRef, PyType, PyTypeRef};
5-
use crate::convert::ToPyObject;
64
use crate::function::{ArgBytesLike, Either, FuncArgs, KwArgs, OptionalArg};
75
use crate::protocol::{BufferDescriptor, BufferMethods, PyBuffer, PyNumberMethods};
86
use crate::stdlib::ctypes::_ctypes::new_simple_type;
@@ -231,10 +229,7 @@ impl PyCData {
231229

232230
#[pyclass(module = "_ctypes", name = "PyCSimpleType", base = PyType)]
233231
#[derive(Debug, PyPayload, Default)]
234-
pub struct PyCSimpleType {
235-
#[allow(dead_code)]
236-
pub stg_info: StgInfo,
237-
}
232+
pub struct PyCSimpleType {}
238233

239234
#[pyclass(flags(BASETYPE), with(AsNumber))]
240235
impl PyCSimpleType {
@@ -747,6 +742,8 @@ impl PyCSimple {
747742
#[pyclassmethod]
748743
fn repeat(cls: PyTypeRef, n: isize, vm: &VirtualMachine) -> PyResult {
749744
use super::_ctypes::get_size;
745+
use super::array::create_array_type_with_stginfo;
746+
750747
if n < 0 {
751748
return Err(vm.new_value_error(format!("Array length must be >= 0, not {n}")));
752749
}
@@ -766,14 +763,14 @@ impl PyCSimple {
766763
std::mem::size_of::<usize>()
767764
};
768765
let total_size = element_size * (n as usize);
769-
let stg_info = super::util::StgInfo::new(total_size, element_size);
770-
Ok(PyCArrayType {
771-
stg_info,
772-
typ: PyRwLock::new(cls.clone().into()),
773-
length: AtomicCell::new(n as usize),
774-
element_size: AtomicCell::new(element_size),
775-
}
776-
.to_pyobject(vm))
766+
let stg_info = super::util::StgInfo::new_array(
767+
total_size,
768+
element_size,
769+
n as usize,
770+
cls.clone().into(),
771+
element_size,
772+
);
773+
create_array_type_with_stginfo(stg_info, vm)
777774
}
778775

779776
#[pyclassmethod]

0 commit comments

Comments
 (0)