Skip to content

Commit 5a2104e

Browse files
youknowoneyouknowone
committed
fix sizeof/alignment
1 parent 9c98687 commit 5a2104e

File tree

1 file changed

+159
-8
lines changed

1 file changed

+159
-8
lines changed

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

Lines changed: 159 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -194,17 +194,96 @@ pub(crate) mod _ctypes {
194194

195195
#[pyfunction(name = "sizeof")]
196196
pub fn size_of(tp: Either<PyTypeRef, PyObjectRef>, vm: &VirtualMachine) -> PyResult<usize> {
197+
use super::array::{PyCArray, PyCArrayType};
198+
use super::structure::PyCStructure;
199+
use num_traits::ToPrimitive;
200+
197201
match tp {
198202
Either::A(type_) if type_.fast_issubclass(PyCSimple::static_type()) => {
199203
let zelf = new_simple_type(Either::B(&type_), vm)?;
200204
Ok(get_size(zelf._type_.as_str()))
201205
}
202-
Either::B(obj) if obj.has_attr("size_of_instances", vm)? => {
203-
let size_of_method = obj.get_attr("size_of_instances", vm)?;
204-
let size_of_return = size_of_method.call(vec![], vm)?;
205-
Ok(usize::try_from_object(vm, size_of_return)?)
206+
Either::A(type_) => {
207+
// Check if it's PyCArrayType (metatype)
208+
if let Some(array_type) = type_.as_object().downcast_ref::<PyCArrayType>() {
209+
let element_type = array_type.inner.typ.read().clone();
210+
let length = array_type.inner.length.load();
211+
let element_size = size_of(Either::A(element_type), vm).unwrap_or(0);
212+
return Ok(element_size * length);
213+
}
214+
// Check if it's a Structure, Union, or Array type
215+
// Try to get _fields_ for Structure/Union
216+
if let Ok(fields_attr) = type_.as_object().get_attr("_fields_", vm) {
217+
// Calculate size from fields
218+
let mut total_size = 0usize;
219+
if let Ok(fields) = fields_attr.try_to_value::<Vec<PyObjectRef>>(vm) {
220+
for field in fields.iter() {
221+
let Some(tuple) = field.downcast_ref::<crate::builtins::PyTuple>()
222+
else {
223+
continue;
224+
};
225+
let Some(field_type) = tuple.get(1).filter(|_| tuple.len() >= 2) else {
226+
continue;
227+
};
228+
// field_type is a type object, so use Either::A if possible
229+
let field_size = if let Ok(ft) =
230+
field_type.clone().downcast::<crate::builtins::PyType>()
231+
{
232+
size_of(Either::A(ft), vm).unwrap_or(0)
233+
} else {
234+
size_of(Either::B(field_type.clone()), vm).unwrap_or(0)
235+
};
236+
total_size += field_size;
237+
}
238+
}
239+
if total_size > 0 {
240+
return Ok(total_size);
241+
}
242+
}
243+
// Try to get _type_ and _length_ for Array type
244+
if let (Ok(type_attr), Ok(length_attr)) = (
245+
type_.as_object().get_attr("_type_", vm),
246+
type_.as_object().get_attr("_length_", vm),
247+
) {
248+
let element_size =
249+
if let Ok(ta) = type_attr.clone().downcast::<crate::builtins::PyType>() {
250+
size_of(Either::A(ta), vm).unwrap_or(0)
251+
} else {
252+
size_of(Either::B(type_attr), vm).unwrap_or(0)
253+
};
254+
let length = length_attr.try_int(vm)?.as_bigint().to_usize().unwrap_or(0);
255+
return Ok(element_size * length);
256+
}
257+
Err(vm.new_type_error("this type has no size"))
258+
}
259+
Either::B(obj) => {
260+
// Handle PyCArrayType (result of c_int * 5)
261+
if let Some(array_type) = obj.downcast_ref::<PyCArrayType>() {
262+
let element_type = array_type.inner.typ.read().clone();
263+
let length = array_type.inner.length.load();
264+
let element_size = size_of(Either::A(element_type), vm).unwrap_or(0);
265+
return Ok(element_size * length);
266+
}
267+
// Handle instances
268+
if let Some(structure) = obj.downcast_ref::<PyCStructure>() {
269+
return Ok(structure.size.load());
270+
}
271+
if let Some(array) = obj.downcast_ref::<PyCArray>() {
272+
return Ok(array.length.load() * array.element_size.load());
273+
}
274+
if let Some(simple) = obj.downcast_ref::<PyCSimple>() {
275+
return Ok(get_size(&simple._type_));
276+
}
277+
// Try size_of_instances method
278+
if obj.has_attr("size_of_instances", vm)? {
279+
let size_of_method = obj.get_attr("size_of_instances", vm)?;
280+
let size_of_return = size_of_method.call(vec![], vm)?;
281+
return Ok(usize::try_from_object(vm, size_of_return)?);
282+
}
283+
// Try to get the class and calculate from there
284+
let class = obj.class();
285+
size_of(Either::A(class.to_owned()), vm)
206286
}
207-
_ => Err(vm.new_type_error("this type has no size")),
208287
}
209288
}
210289

@@ -365,9 +444,81 @@ pub(crate) mod _ctypes {
365444
}
366445

367446
#[pyfunction]
368-
fn alignment(_args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> {
369-
// TODO: RUSTPYTHON
370-
Err(vm.new_value_error("not implemented"))
447+
fn alignment(tp: Either<PyTypeRef, PyObjectRef>, vm: &VirtualMachine) -> PyResult<usize> {
448+
use super::array::{PyCArray, PyCArrayType};
449+
use super::structure::PyCStructure;
450+
451+
// Get alignment based on the type
452+
// For simple types, alignment equals size
453+
// For composite types, alignment is the max alignment of members
454+
match tp {
455+
Either::A(type_) if type_.fast_issubclass(PyCSimple::static_type()) => {
456+
let zelf = new_simple_type(Either::B(&type_), vm)?;
457+
Ok(get_size(zelf._type_.as_str()))
458+
}
459+
Either::A(type_) => {
460+
// Check if it's PyCArrayType (metatype)
461+
if let Some(array_type) = type_.as_object().downcast_ref::<PyCArrayType>() {
462+
let element_type = array_type.inner.typ.read().clone();
463+
return alignment(Either::A(element_type), vm);
464+
}
465+
// Check for _fields_ (Structure/Union)
466+
if let Ok(fields_attr) = type_.as_object().get_attr("_fields_", vm) {
467+
let mut max_align = 1usize;
468+
if let Ok(fields) = fields_attr.try_to_value::<Vec<PyObjectRef>>(vm) {
469+
for field in fields.iter() {
470+
let Some(tuple) = field.downcast_ref::<crate::builtins::PyTuple>()
471+
else {
472+
continue;
473+
};
474+
let Some(field_type) = tuple.get(1).filter(|_| tuple.len() >= 2) else {
475+
continue;
476+
};
477+
let field_align = if let Ok(ft) =
478+
field_type.clone().downcast::<crate::builtins::PyType>()
479+
{
480+
alignment(Either::A(ft), vm).unwrap_or(1)
481+
} else {
482+
alignment(Either::B(field_type.clone()), vm).unwrap_or(1)
483+
};
484+
max_align = max_align.max(field_align);
485+
}
486+
}
487+
return Ok(max_align);
488+
}
489+
// Check for _type_ (Array)
490+
if let Ok(type_attr) = type_.as_object().get_attr("_type_", vm) {
491+
return if let Ok(ta) = type_attr.clone().downcast::<crate::builtins::PyType>() {
492+
alignment(Either::A(ta), vm)
493+
} else {
494+
alignment(Either::B(type_attr), vm)
495+
};
496+
}
497+
// Default alignment
498+
Ok(1)
499+
}
500+
Either::B(obj) => {
501+
// Handle PyCArrayType (result of c_int * 5)
502+
if let Some(array_type) = obj.downcast_ref::<PyCArrayType>() {
503+
let element_type = array_type.inner.typ.read().clone();
504+
return alignment(Either::A(element_type), vm);
505+
}
506+
// Handle instances
507+
if let Some(_structure) = obj.downcast_ref::<PyCStructure>() {
508+
// Get alignment from class
509+
return alignment(Either::A(obj.class().to_owned()), vm);
510+
}
511+
if let Some(array) = obj.downcast_ref::<PyCArray>() {
512+
return Ok(array.element_size.load().min(8)); // Max alignment is typically 8
513+
}
514+
if let Some(simple) = obj.downcast_ref::<PyCSimple>() {
515+
return Ok(get_size(&simple._type_));
516+
}
517+
// Try to get the class and calculate from there
518+
let class = obj.class();
519+
alignment(Either::A(class.to_owned()), vm)
520+
}
521+
}
371522
}
372523

373524
#[pyfunction]

0 commit comments

Comments
 (0)