Skip to content

Commit 274519c

Browse files
youknowoneyouknowone
committed
fix sizeof/alignment
1 parent bea4eec commit 274519c

9 files changed

Lines changed: 461 additions & 179 deletions

File tree

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

Lines changed: 135 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,13 @@ pub(crate) mod pointer;
99
pub(crate) mod structure;
1010
pub(crate) mod thunk;
1111
pub(crate) mod union;
12+
pub(crate) mod util;
1213

1314
use crate::builtins::PyModule;
1415
use crate::class::PyClassImpl;
1516
use crate::{Py, PyRef, VirtualMachine};
1617

17-
pub use crate::stdlib::ctypes::base::{CDataObject, PyCData, PyCSimple, PyCSimpleType};
18+
pub use crate::stdlib::ctypes::base::{PyCData, PyCSimple, PyCSimpleType};
1819

1920
pub fn extend_module_nodes(vm: &VirtualMachine, module: &Py<PyModule>) {
2021
let ctx = &vm.ctx;
@@ -52,7 +53,7 @@ pub(crate) mod _ctypes {
5253
use crate::convert::ToPyObject;
5354
use crate::function::{Either, FuncArgs, OptionalArg};
5455
use crate::stdlib::ctypes::library;
55-
use crate::{AsObject, PyObjectRef, PyPayload, PyResult, TryFromObject, VirtualMachine};
56+
use crate::{AsObject, PyObjectRef, PyPayload, PyResult, VirtualMachine};
5657
use crossbeam_utils::atomic::AtomicCell;
5758
use std::ffi::{
5859
c_double, c_float, c_int, c_long, c_longlong, c_schar, c_short, c_uchar, c_uint, c_ulong,
@@ -159,6 +160,11 @@ pub(crate) mod _ctypes {
159160
}
160161
}
161162

163+
/// Get alignment for a simple type - for C types, alignment equals size
164+
pub fn get_align(ty: &str) -> usize {
165+
get_size(ty)
166+
}
167+
162168
/// Get the size of a ctypes type from its type object
163169
#[allow(dead_code)]
164170
pub fn get_size_from_type(cls: &PyTypeRef, vm: &VirtualMachine) -> PyResult<usize> {
@@ -366,7 +372,7 @@ pub(crate) mod _ctypes {
366372
Ok(PyCSimple {
367373
_type_: tp_str,
368374
value: AtomicCell::new(vm.ctx.none()),
369-
cdata: rustpython_common::lock::PyRwLock::new(CDataObject::new(size)),
375+
cdata: rustpython_common::lock::PyRwLock::new(CDataObject::from_bytes(vec![0u8; size], None)),
370376
})
371377
}
372378
} else {
@@ -378,95 +384,52 @@ pub(crate) mod _ctypes {
378384
}
379385

380386
/// Get the size of a ctypes type or instance
381-
///
382-
/// This function accepts:
383-
/// - A ctypes type (e.g., c_int, Structure subclass)
384-
/// - A ctypes instance (e.g., c_int(5))
385387
#[pyfunction(name = "sizeof")]
386388
pub fn size_of(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<usize> {
387-
use super::array::PyCArray;
389+
use super::array::{PyCArray, PyCArrayType};
388390
use super::pointer::PyCPointer;
389391
use super::structure::PyCStructure;
390-
use super::union::PyCUnion;
391-
392-
// Check if obj is a type
393-
if let Ok(type_ref) = obj.clone().downcast::<crate::builtins::PyType>() {
394-
// It's a type - check what kind of ctypes type it is
395-
396-
// Simple types (c_int, c_char, etc.)
397-
if type_ref.fast_issubclass(PyCSimple::static_type()) {
398-
let zelf = new_simple_type(Either::B(&type_ref), vm)?;
399-
return Ok(get_size(zelf._type_.as_str()));
400-
}
401-
402-
// Array types
403-
if type_ref.fast_issubclass(PyCArray::static_type()) {
404-
// Get _length_ and element size
405-
if let Ok(length) = type_ref.as_object().get_attr("_length_", vm) {
406-
let length = usize::try_from_object(vm, length)?;
407-
if let Ok(elem_type) = type_ref.as_object().get_attr("_type_", vm) {
408-
let elem_size = size_of(elem_type, vm)?;
409-
return Ok(length * elem_size);
410-
}
411-
}
412-
}
413-
414-
// Structure types - check for size_of_instances method
415-
if type_ref.fast_issubclass(PyCStructure::static_type())
416-
|| type_ref.fast_issubclass(PyCUnion::static_type())
417-
{
418-
if let Ok(size_method) = type_ref.as_object().get_attr("size_of_instances", vm) {
419-
let size = size_method.call(vec![], vm)?;
420-
return Ok(usize::try_from_object(vm, size)?);
421-
}
422-
}
423-
424-
// Pointer types
425-
if type_ref.fast_issubclass(PyCPointer::static_type()) {
426-
return Ok(std::mem::size_of::<usize>());
427-
}
392+
use super::union::{PyCUnion, PyCUnionType};
428393

429-
// Check for size_of_instances as a fallback
430-
if let Ok(size_method) = type_ref.as_object().get_attr("size_of_instances", vm) {
431-
let size = size_method.call(vec![], vm)?;
432-
return Ok(usize::try_from_object(vm, size)?);
394+
// 1. Instances with stg_info
395+
if obj.fast_isinstance(PyCArray::static_type()) {
396+
// Get stg_info from the type
397+
if let Some(type_obj) = obj.class().as_object().downcast_ref::<PyCArrayType>() {
398+
return Ok(type_obj.stg_info.size);
433399
}
434-
435-
return Err(vm.new_type_error("this type has no size"));
436-
}
437-
438-
// It's an instance - get size from instance or its class
439-
440-
// Simple type instance
441-
if let Ok(simple) = obj.clone().downcast::<PyCSimple>() {
442-
return Ok(get_size(simple._type_.as_str()));
443-
}
444-
445-
// Array instance
446-
if let Ok(array) = obj.clone().downcast::<PyCArray>() {
447-
return Ok(array.cdata.read().size());
448400
}
449-
450-
// Structure instance
451-
if let Ok(structure) = obj.clone().downcast::<PyCStructure>() {
401+
if let Some(structure) = obj.downcast_ref::<PyCStructure>() {
452402
return Ok(structure.cdata.read().size());
453403
}
454-
455-
// Union instance
456-
if let Ok(union) = obj.clone().downcast::<PyCUnion>() {
457-
return Ok(union.cdata.read().size());
404+
if obj.fast_isinstance(PyCUnion::static_type()) {
405+
// Get stg_info from the type
406+
if let Some(type_obj) = obj.class().as_object().downcast_ref::<PyCUnionType>() {
407+
return Ok(type_obj.stg_info.size);
408+
}
409+
}
410+
if let Some(simple) = obj.downcast_ref::<PyCSimple>() {
411+
return Ok(simple.cdata.read().size());
458412
}
459-
460-
// Pointer instance
461413
if obj.fast_isinstance(PyCPointer::static_type()) {
462414
return Ok(std::mem::size_of::<usize>());
463415
}
464416

465-
// Check if the object has size_of_instances method (for custom types)
466-
if obj.has_attr("size_of_instances", vm)? {
467-
let size_method = obj.get_attr("size_of_instances", vm)?;
468-
let size = size_method.call(vec![], vm)?;
469-
return Ok(usize::try_from_object(vm, size)?);
417+
// 2. Types (metatypes with stg_info)
418+
if let Some(array_type) = obj.downcast_ref::<PyCArrayType>() {
419+
return Ok(array_type.stg_info.size);
420+
}
421+
422+
// 3. Type objects
423+
if let Ok(type_ref) = obj.clone().downcast::<crate::builtins::PyType>() {
424+
// Simple types (c_int, c_char, etc.)
425+
if type_ref.fast_issubclass(PyCSimple::static_type()) {
426+
let instance = new_simple_type(Either::B(&type_ref), vm)?;
427+
return Ok(get_size(&instance._type_));
428+
}
429+
// Pointer types
430+
if type_ref.fast_issubclass(PyCPointer::static_type()) {
431+
return Ok(std::mem::size_of::<usize>());
432+
}
470433
}
471434

472435
Err(vm.new_type_error("this type has no size"))
@@ -630,9 +593,100 @@ pub(crate) mod _ctypes {
630593
}
631594

632595
#[pyfunction]
633-
fn alignment(_args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> {
634-
// TODO: RUSTPYTHON
635-
Err(vm.new_value_error("not implemented"))
596+
fn alignment(tp: Either<PyTypeRef, PyObjectRef>, vm: &VirtualMachine) -> PyResult<usize> {
597+
use super::array::{PyCArray, PyCArrayType};
598+
use super::base::PyCSimpleType;
599+
use super::pointer::PyCPointer;
600+
use super::structure::PyCStructure;
601+
use super::union::PyCUnion;
602+
603+
let obj = match &tp {
604+
Either::A(t) => t.as_object(),
605+
Either::B(o) => o.as_ref(),
606+
};
607+
608+
// Try to get alignment from stg_info directly (for instances)
609+
if let Some(array_type) = obj.downcast_ref::<PyCArrayType>() {
610+
return Ok(array_type.stg_info.align);
611+
}
612+
if obj.fast_isinstance(PyCSimple::static_type()) {
613+
// Get stg_info from the type by reading _type_ attribute
614+
let cls = obj.class().to_owned();
615+
let stg_info = PyCSimpleType::get_stg_info(&cls, vm);
616+
return Ok(stg_info.align);
617+
}
618+
if obj.fast_isinstance(PyCArray::static_type()) {
619+
// Get stg_info from the type
620+
if let Some(type_obj) = obj.class().as_object().downcast_ref::<PyCArrayType>() {
621+
return Ok(type_obj.stg_info.align);
622+
}
623+
}
624+
if obj.fast_isinstance(PyCStructure::static_type()) {
625+
// Calculate alignment from _fields_
626+
let cls = obj.class();
627+
return alignment(Either::A(cls.to_owned()), vm);
628+
}
629+
if obj.fast_isinstance(PyCPointer::static_type()) {
630+
// Pointer alignment is always pointer size
631+
return Ok(std::mem::align_of::<usize>());
632+
}
633+
if obj.fast_isinstance(PyCUnion::static_type()) {
634+
// Calculate alignment from _fields_
635+
let cls = obj.class();
636+
return alignment(Either::A(cls.to_owned()), vm);
637+
}
638+
639+
// Get the type object to check
640+
let type_obj: PyObjectRef = match &tp {
641+
Either::A(t) => t.clone().into(),
642+
Either::B(obj) => obj.class().to_owned().into(),
643+
};
644+
645+
// For type objects, try to get alignment from _type_ attribute
646+
if let Ok(type_attr) = type_obj.get_attr("_type_", vm) {
647+
// Array/Pointer: _type_ is the element type (a PyType)
648+
if let Ok(elem_type) = type_attr.clone().downcast::<crate::builtins::PyType>() {
649+
return alignment(Either::A(elem_type), vm);
650+
}
651+
// Simple type: _type_ is a single character string
652+
if let Ok(s) = type_attr.str(vm) {
653+
let ty = s.to_string();
654+
if ty.len() == 1 && SIMPLE_TYPE_CHARS.contains(ty.as_str()) {
655+
return Ok(get_align(&ty));
656+
}
657+
}
658+
}
659+
660+
// Structure/Union: max alignment of fields
661+
if let Ok(fields_attr) = type_obj.get_attr("_fields_", vm)
662+
&& let Ok(fields) = fields_attr.try_to_value::<Vec<PyObjectRef>>(vm)
663+
{
664+
let mut max_align = 1usize;
665+
for field in fields.iter() {
666+
if let Some(tuple) = field.downcast_ref::<crate::builtins::PyTuple>()
667+
&& let Some(field_type) = tuple.get(1)
668+
{
669+
let align =
670+
if let Ok(ft) = field_type.clone().downcast::<crate::builtins::PyType>() {
671+
alignment(Either::A(ft), vm).unwrap_or(1)
672+
} else {
673+
1
674+
};
675+
max_align = max_align.max(align);
676+
}
677+
}
678+
return Ok(max_align);
679+
}
680+
681+
// For instances, delegate to their class
682+
if let Either::B(obj) = &tp
683+
&& !obj.class().is(vm.ctx.types.type_type.as_ref())
684+
{
685+
return alignment(Either::A(obj.class().to_owned()), vm);
686+
}
687+
688+
// No alignment info found
689+
Err(vm.new_type_error("no alignment info"))
636690
}
637691

638692
#[pyfunction]

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

Lines changed: 24 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use crate::protocol::{
66
BufferDescriptor, BufferMethods, PyBuffer, PyNumberMethods, PySequenceMethods,
77
};
88
use crate::stdlib::ctypes::base::CDataObject;
9+
use crate::stdlib::ctypes::util::StgInfo;
910
use crate::types::{AsBuffer, AsNumber, AsSequence, Callable};
1011
use crate::{AsObject, Py, PyObjectRef, PyPayload};
1112
use crate::{
@@ -22,13 +23,17 @@ use rustpython_vm::stdlib::ctypes::base::PyCData;
2223
#[pyclass(name = "PyCArrayType", base = PyType, module = "_ctypes")]
2324
#[derive(PyPayload)]
2425
pub struct PyCArrayType {
25-
pub(super) inner: PyCArray,
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>,
2630
}
2731

2832
impl std::fmt::Debug for PyCArrayType {
2933
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3034
f.debug_struct("PyCArrayType")
31-
.field("inner", &self.inner)
35+
.field("typ", &self.typ)
36+
.field("length", &self.length)
3237
.finish()
3338
}
3439
}
@@ -37,9 +42,9 @@ impl Callable for PyCArrayType {
3742
type Args = FuncArgs;
3843
fn call(zelf: &Py<Self>, args: Self::Args, vm: &VirtualMachine) -> PyResult {
3944
// Create an instance of the array
40-
let element_type = zelf.inner.typ.read().clone();
41-
let length = zelf.inner.length.load();
42-
let element_size = zelf.inner.element_size.load();
45+
let element_type = zelf.typ.read().clone();
46+
let length = zelf.length.load();
47+
let element_size = zelf.element_size.load();
4348
let total_size = element_size * length;
4449
let mut buffer = vec![0u8; total_size];
4550

@@ -79,12 +84,12 @@ impl Constructor for PyCArrayType {
7984
impl PyCArrayType {
8085
#[pygetset(name = "_type_")]
8186
fn typ(&self) -> PyObjectRef {
82-
self.inner.typ.read().clone()
87+
self.typ.read().clone()
8388
}
8489

8590
#[pygetset(name = "_length_")]
8691
fn length(&self) -> usize {
87-
self.inner.length.load()
92+
self.length.load()
8893
}
8994

9095
#[pymethod]
@@ -95,22 +100,22 @@ impl PyCArrayType {
95100
// Create a nested array type: (inner_type * inner_length) * n
96101
// The new array has n elements, each element is the current array type
97102
// e.g., (c_int * 5) * 3 = Array of 3 elements, each is (c_int * 5)
98-
let inner_length = zelf.inner.length.load();
99-
let inner_element_size = zelf.inner.element_size.load();
103+
let inner_length = zelf.length.load();
104+
let inner_element_size = zelf.element_size.load();
100105

101106
// The element type of the new array is the current array type itself
102107
let current_array_type: PyObjectRef = zelf.as_object().to_owned();
103108

104109
// Element size is the total size of the inner array
105110
let new_element_size = inner_length * inner_element_size;
111+
let total_size = new_element_size * (n as usize);
112+
let stg_info = StgInfo::new(total_size, inner_element_size);
106113

107114
Ok(PyCArrayType {
108-
inner: PyCArray {
109-
typ: PyRwLock::new(current_array_type),
110-
length: AtomicCell::new(n as usize),
111-
element_size: AtomicCell::new(new_element_size),
112-
cdata: PyRwLock::new(CDataObject::new(0)),
113-
},
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),
114119
}
115120
.to_pyobject(vm))
116121
}
@@ -164,8 +169,8 @@ impl PyCArrayType {
164169
};
165170

166171
// Get size from the array type
167-
let element_type = zelf.inner.typ.read().clone();
168-
let length = zelf.inner.length.load();
172+
let element_type = zelf.typ.read().clone();
173+
let length = zelf.length.load();
169174
let element_size = size_of(element_type.clone(), vm)?;
170175
let total_size = element_size * length;
171176

@@ -184,7 +189,7 @@ impl PyCArrayType {
184189
typ: PyRwLock::new(element_type),
185190
length: AtomicCell::new(length),
186191
element_size: AtomicCell::new(element_size),
187-
cdata: PyRwLock::new(CDataObject::from_bytes(data)),
192+
cdata: PyRwLock::new(CDataObject::from_bytes(data, None)),
188193
}
189194
.into_pyobject(vm);
190195

@@ -729,7 +734,7 @@ impl PyCArray {
729734
typ: PyRwLock::new(element_type.into()),
730735
length: AtomicCell::new(length),
731736
element_size: AtomicCell::new(element_size),
732-
cdata: PyRwLock::new(CDataObject::from_bytes(data)),
737+
cdata: PyRwLock::new(CDataObject::from_bytes(data, None)),
733738
}
734739
.into_pyobject(vm);
735740

0 commit comments

Comments
 (0)