Use sequence repetition for ctypes array types (#9185)

Assisted-by: Codex:GPT-6

Co-authored-by: Jeong, YunWon <69878+youknowone@users.noreply.github.com>

crates/vm/src/stdlib/_ctypes/array.rs

@@ -10,8 +10,8 @@ use crate::{
     },
     class::StaticType,
     function::{ArgBytesLike, FuncArgs, PySetterValue},
-    protocol::{BufferDescriptor, PyBuffer, PyMappingMethods, PyNumberMethods, PySequenceMethods},
-    types::{AsBuffer, AsMapping, AsNumber, AsSequence, Constructor, Initializer},
+    protocol::{BufferDescriptor, PyBuffer, PyMappingMethods, PySequenceMethods},
+    types::{AsBuffer, AsMapping, AsSequence, Constructor, Initializer},
 };
 use alloc::borrow::Cow;
 use num_traits::{Signed, ToPrimitive};
@@ -44,6 +44,27 @@ fn get_size_from_format(fmt: &str) -> usize {
     code.map_or(1, |c| simple_type_size(&c).unwrap_or(1))
 }
 
+pub(super) static CTYPE_SEQUENCE_METHODS: PySequenceMethods = PySequenceMethods {
+    repeat: Some(|seq, length, vm| repeat_type(seq.obj, length, vm)),
+    ..PySequenceMethods::NOT_IMPLEMENTED
+};
+
+pub(super) fn repeat_type(ctype: &PyObject, length: isize, vm: &VirtualMachine) -> PyResult {
+    if length < 0 {
+        return Err(vm.new_value_error(format!("Array length must be >= 0, not {length}")));
+    }
+    let cls = ctype
+        .downcast_ref::<PyType>()
+        .ok_or_else(|| vm.new_type_error("expected type"))?;
+    if let Some(info) = cls.stg_info_opt()
+        && info.size != 0
+        && length as usize > isize::MAX as usize / info.size
+    {
+        return Err(vm.new_overflow_error("array too large"));
+    }
+    array_type_from_ctype(ctype.to_owned(), length as usize, vm)
+}
+
 /// Creates array type for (element_type, length)
 /// Uses _array_type_cache to ensure identical calls return the same type object
 pub(super) fn array_type_from_ctype(
@@ -305,7 +326,7 @@ impl Initializer for PyCArrayType {
     }
 }
 
-#[pyclass(flags(IMMUTABLETYPE), with(Initializer, AsNumber))]
+#[pyclass(flags(IMMUTABLETYPE), with(Initializer, AsSequence))]
 impl PyCArrayType {
     #[pygetset(name = "__pointer_type__")]
     fn pointer_type(zelf: PyTypeRef, vm: &VirtualMachine) -> PyResult {
@@ -350,41 +371,9 @@ impl PyCArrayType {
     }
 }
 
-impl AsNumber for PyCArrayType {
-    fn as_number() -> &'static PyNumberMethods {
-        static AS_NUMBER: PyNumberMethods = PyNumberMethods {
-            multiply: Some(|a, b, vm| {
-                // a is a type object whose metaclass is PyCArrayType (e.g., Array_5)
-                let n = b
-                    .try_index(vm)?
-                    .as_bigint()
-                    .to_isize()
-                    .ok_or_else(|| vm.new_overflow_error("array size too large"))?;
-
-                if n < 0 {
-                    return Err(vm.new_value_error(format!("Array length must be >= 0, not {n}")));
-                }
-
-                // Check for overflow before creating the new array type
-                let zelf_type = a
-                    .downcast_ref::<PyType>()
-                    .ok_or_else(|| vm.new_type_error("Expected type"))?;
-
-                if let Some(stg_info) = zelf_type.stg_info_opt() {
-                    let current_size = stg_info.size;
-                    // Check if current_size * n would overflow
-                    if current_size != 0 && (n as usize) > isize::MAX as usize / current_size {
-                        return Err(vm.new_overflow_error("array too large"));
-                    }
-                }
-
-                // Use cached array type creation
-                // The element type of the new array is the current array type itself
-                array_type_from_ctype(a.to_owned(), n as usize, vm)
-            }),
-            ..PyNumberMethods::NOT_IMPLEMENTED
-        };
-        &AS_NUMBER
+impl AsSequence for PyCArrayType {
+    fn as_sequence() -> &'static PySequenceMethods {
+        &super::array::CTYPE_SEQUENCE_METHODS
     }
 }
 

crates/vm/src/stdlib/_ctypes/function.rs

@@ -12,8 +12,8 @@ use crate::{
     builtins::{PyBytes, PyDict, PyInt, PyStr, PyTuple, PyType, PyTypeRef},
     class::StaticType,
     function::FuncArgs,
-    protocol::{BufferDescriptor, PyBuffer, PyNumberMethods},
-    types::{AsBuffer, AsNumber, Callable, Constructor, Initializer, Representable},
+    protocol::{BufferDescriptor, PyBuffer, PyNumberMethods, PySequenceMethods},
+    types::{AsBuffer, AsNumber, AsSequence, Callable, Constructor, Initializer, Representable},
     vm::thread::with_current_vm,
 };
 use alloc::borrow::Cow;
@@ -331,7 +331,7 @@ impl Initializer for PyCFuncPtrType {
     }
 }
 
-#[pyclass(flags(IMMUTABLETYPE), with(Initializer))]
+#[pyclass(flags(IMMUTABLETYPE), with(AsSequence, Initializer))]
 impl PyCFuncPtrType {
     #[pygetset(name = "__pointer_type__")]
     fn pointer_type(zelf: PyTypeRef, vm: &VirtualMachine) -> PyResult {
@@ -2032,6 +2032,12 @@ impl Py<PyCThunk> {
     }
 }
 
+impl AsSequence for PyCFuncPtrType {
+    fn as_sequence() -> &'static PySequenceMethods {
+        &super::array::CTYPE_SEQUENCE_METHODS
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/vm/src/stdlib/_ctypes/pointer.rs

@@ -1,8 +1,10 @@
 use super::base::CDATA_BUFFER_METHODS;
 use super::{PyCArray, PyCData, PyCSimple, PyCStructure, StgInfo, StgInfoFlags};
 use crate::atomic_func;
-use crate::protocol::{BufferDescriptor, PyBuffer, PyMappingMethods, PyNumberMethods};
-use crate::types::{AsBuffer, AsMapping, AsNumber, Constructor, Initializer};
+use crate::protocol::{
+    BufferDescriptor, PyBuffer, PyMappingMethods, PyNumberMethods, PySequenceMethods,
+};
+use crate::types::{AsBuffer, AsMapping, AsNumber, AsSequence, Constructor, Initializer};
 use crate::{
     AsObject, Py, PyObject, PyObjectRef, PyPayload, PyResult, VirtualMachine,
     builtins::{PyBytes, PyInt, PyList, PySlice, PyStr, PyType, PyTypeRef},
@@ -87,7 +89,7 @@ impl Initializer for PyCPointerType {
     }
 }
 
-#[pyclass(flags(BASETYPE, IMMUTABLETYPE), with(AsNumber, Initializer))]
+#[pyclass(flags(BASETYPE, IMMUTABLETYPE), with(AsSequence, Initializer))]
 impl PyCPointerType {
     #[pygetset(name = "__pointer_type__")]
     fn pointer_type(zelf: PyTypeRef, vm: &VirtualMachine) -> PyResult {
@@ -175,16 +177,6 @@ impl PyCPointerType {
         )))
     }
 
-    fn __mul__(cls: PyTypeRef, n: isize, vm: &VirtualMachine) -> PyResult {
-        use super::array::array_type_from_ctype;
-
-        if n < 0 {
-            return Err(vm.new_value_error(format!("Array length must be >= 0, not {n}")));
-        }
-        // Use cached array type creation
-        array_type_from_ctype(cls.into(), n as usize, vm)
-    }
-
     // PyCPointerType_set_type: Complete an incomplete pointer type
     #[pymethod]
     fn set_type(zelf: PyTypeRef, typ: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
@@ -229,23 +221,9 @@ impl PyCPointerType {
     }
 }
 
-impl AsNumber for PyCPointerType {
-    fn as_number() -> &'static PyNumberMethods {
-        static AS_NUMBER: PyNumberMethods = PyNumberMethods {
-            multiply: Some(|a, b, vm| {
-                let cls = a
-                    .downcast_ref::<PyType>()
-                    .ok_or_else(|| vm.new_type_error("expected type"))?;
-                let n = b
-                    .try_index(vm)?
-                    .as_bigint()
-                    .to_isize()
-                    .ok_or_else(|| vm.new_overflow_error("array size too large"))?;
-                PyCPointerType::__mul__(cls.to_owned(), n, vm)
-            }),
-            ..PyNumberMethods::NOT_IMPLEMENTED
-        };
-        &AS_NUMBER
+impl AsSequence for PyCPointerType {
+    fn as_sequence() -> &'static PySequenceMethods {
+        &super::array::CTYPE_SEQUENCE_METHODS
     }
 }
 

crates/vm/src/stdlib/_ctypes/simple.rs

@@ -8,8 +8,8 @@ use super::pointer::PyCPointer;
 use crate::builtins::{PyByteArray, PyBytes, PyInt, PyNone, PyStr, PyType, PyTypeRef};
 use crate::convert::ToPyObject;
 use crate::function::{Either, FuncArgs, OptionalArg};
-use crate::protocol::{BufferDescriptor, PyBuffer, PyNumberMethods};
-use crate::types::{AsBuffer, AsNumber, Constructor, Initializer, Representable};
+use crate::protocol::{BufferDescriptor, PyBuffer, PyNumberMethods, PySequenceMethods};
+use crate::types::{AsBuffer, AsNumber, AsSequence, Constructor, Initializer, Representable};
 use crate::{
     AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
     class::PyClassDef, set_attrs,
@@ -226,7 +226,7 @@ fn set_primitive(_type_: &str, value: &PyObject, vm: &VirtualMachine) -> PyResul
 #[repr(transparent)]
 pub(crate) struct PyCSimpleType(PyType);
 
-#[pyclass(flags(BASETYPE), with(AsNumber, Initializer))]
+#[pyclass(flags(BASETYPE), with(AsSequence, Initializer))]
 impl PyCSimpleType {
     #[pygetset(name = "__pointer_type__")]
     fn pointer_type(zelf: PyTypeRef, vm: &VirtualMachine) -> PyResult {
@@ -494,31 +494,11 @@ impl PyCSimpleType {
             _ => Err(vm.new_type_error("wrong type")),
         }
     }
-
-    fn __mul__(cls: PyTypeRef, n: isize, vm: &VirtualMachine) -> PyResult {
-        Py::<PyCSimple>::repeat(cls, n, vm)
-    }
 }
 
-impl AsNumber for PyCSimpleType {
-    fn as_number() -> &'static PyNumberMethods {
-        static AS_NUMBER: PyNumberMethods = PyNumberMethods {
-            multiply: Some(|a, b, vm| {
-                // a is a PyCSimpleType instance (type object like c_char)
-                // b is int (array size)
-                let cls = a
-                    .downcast_ref::<PyType>()
-                    .ok_or_else(|| vm.new_type_error("expected type"))?;
-                let n = b
-                    .try_index(vm)?
-                    .as_bigint()
-                    .to_isize()
-                    .ok_or_else(|| vm.new_overflow_error("array size too large"))?;
-                Py::<PyCSimple>::repeat(cls.to_owned(), n, vm)
-            }),
-            ..PyNumberMethods::NOT_IMPLEMENTED
-        };
-        &AS_NUMBER
+impl AsSequence for PyCSimpleType {
+    fn as_sequence() -> &'static PySequenceMethods {
+        &super::array::CTYPE_SEQUENCE_METHODS
     }
 }
 
@@ -1229,13 +1209,7 @@ impl Py<PyCSimple> {
 
     #[pyclassmethod]
     fn repeat(cls: PyTypeRef, n: isize, vm: &VirtualMachine) -> PyResult {
-        use super::array::array_type_from_ctype;
-
-        if n < 0 {
-            return Err(vm.new_value_error(format!("Array length must be >= 0, not {n}")));
-        }
-        // Use cached array type creation
-        array_type_from_ctype(cls.into(), n as usize, vm)
+        super::array::repeat_type(cls.as_object(), n, vm)
     }
 
     // Simple_from_outparm - convert output parameter back to Python value

crates/vm/src/stdlib/_ctypes/structure.rs

@@ -2,9 +2,9 @@ use super::base::{CDATA_BUFFER_METHODS, PyCData, PyCField, StgInfo, StgInfoFlags
 use crate::builtins::{PyList, PyStr, PyTuple, PyType, PyTypeRef, PyUtf8Str};
 use crate::convert::ToPyObject;
 use crate::function::{ArgStrictInt, FuncArgs, OptionalArg, PySetterValue};
-use crate::protocol::{BufferDescriptor, PyBuffer, PyNumberMethods};
+use crate::protocol::{BufferDescriptor, PyBuffer, PySequenceMethods};
 use crate::stdlib::_warnings;
-use crate::types::{AsBuffer, AsNumber, Constructor, Initializer, SetAttr};
+use crate::types::{AsBuffer, AsSequence, Constructor, Initializer, SetAttr};
 use crate::{AsObject, Py, PyObject, PyObjectRef, PyPayload, PyResult, VirtualMachine};
 use alloc::borrow::Cow;
 use core::fmt::Debug;
@@ -129,7 +129,7 @@ impl Initializer for PyCStructType {
     }
 }
 
-#[pyclass(flags(BASETYPE), with(AsNumber, Constructor, Initializer, SetAttr))]
+#[pyclass(flags(BASETYPE), with(AsSequence, Constructor, Initializer, SetAttr))]
 impl PyCStructType {
     #[pygetset(name = "__pointer_type__")]
     fn pointer_type(zelf: PyTypeRef, vm: &VirtualMachine) -> PyResult {
@@ -565,35 +565,11 @@ impl PyCStructType {
 
         Ok(())
     }
-
-    fn __mul__(cls: PyTypeRef, n: isize, vm: &VirtualMachine) -> PyResult {
-        use super::array::array_type_from_ctype;
-
-        if n < 0 {
-            return Err(vm.new_value_error(format!("Array length must be >= 0, not {n}")));
-        }
-        // Use cached array type creation
-        array_type_from_ctype(cls.into(), n as usize, vm)
-    }
 }
 
-impl AsNumber for PyCStructType {
-    fn as_number() -> &'static PyNumberMethods {
-        static AS_NUMBER: PyNumberMethods = PyNumberMethods {
-            multiply: Some(|a, b, vm| {
-                let cls = a
-                    .downcast_ref::<PyType>()
-                    .ok_or_else(|| vm.new_type_error("expected type"))?;
-                let n = b
-                    .try_index(vm)?
-                    .as_bigint()
-                    .to_isize()
-                    .ok_or_else(|| vm.new_overflow_error("array size too large"))?;
-                PyCStructType::__mul__(cls.to_owned(), n, vm)
-            }),
-            ..PyNumberMethods::NOT_IMPLEMENTED
-        };
-        &AS_NUMBER
+impl AsSequence for PyCStructType {
+    fn as_sequence() -> &'static PySequenceMethods {
+        &super::array::CTYPE_SEQUENCE_METHODS
     }
 }
 

crates/vm/src/stdlib/_ctypes/union.rs

@@ -3,9 +3,9 @@ use super::{PyCData, PyCField, StgInfo};
 use crate::builtins::{PyList, PyStr, PyTuple, PyType, PyTypeRef, PyUtf8Str};
 use crate::convert::ToPyObject;
 use crate::function::{ArgBytesLike, ArgStrictInt, FuncArgs, OptionalArg, PySetterValue};
-use crate::protocol::{BufferDescriptor, PyBuffer};
+use crate::protocol::{BufferDescriptor, PyBuffer, PySequenceMethods};
 use crate::stdlib::_warnings;
-use crate::types::{AsBuffer, Constructor, Initializer, SetAttr};
+use crate::types::{AsBuffer, AsSequence, Constructor, Initializer, SetAttr};
 use crate::{AsObject, Py, PyObject, PyObjectRef, PyPayload, PyResult, VirtualMachine};
 use alloc::borrow::Cow;
 use num_traits::ToPrimitive;
@@ -356,7 +356,7 @@ impl PyCUnionType {
     }
 }
 
-#[pyclass(flags(BASETYPE), with(Constructor, Initializer, SetAttr))]
+#[pyclass(flags(BASETYPE), with(AsSequence, Constructor, Initializer, SetAttr))]
 impl PyCUnionType {
     #[pygetset(name = "__pointer_type__")]
     fn pointer_type(zelf: PyTypeRef, vm: &VirtualMachine) -> PyResult {
@@ -697,3 +697,9 @@ impl AsBuffer for PyCUnion {
         Ok(buf)
     }
 }
+
+impl AsSequence for PyCUnionType {
+    fn as_sequence() -> &'static PySequenceMethods {
+        &super::array::CTYPE_SEQUENCE_METHODS
+    }
+}

extra_tests/snippets/stdlib_ctypes.py

@@ -470,6 +470,42 @@ gc.collect()
 assert roundtrip(41) == 42
 
 
+# All ctypes metaclasses use sequence repetition in either direction.
+class RepeatStruct(ctypes.Structure):
+    _fields_ = [("value", ctypes.c_int)]
+
+
+class RepeatUnion(ctypes.Union):
+    _fields_ = [("value", ctypes.c_int)]
+
+
+for element_type in (
+    ctypes.c_int,
+    ctypes.c_int * 1,
+    ctypes.POINTER(ctypes.c_int),
+    RepeatStruct,
+    RepeatUnion,
+    Unary,
+):
+    assert 2 * element_type is element_type * 2
+
+
+class ArrayLength:
+    def __index__(self):
+        return 2
+
+
+assert ArrayLength() * ctypes.c_int is ctypes.c_int * 2
+
+
+class ReflectedLength(ArrayLength):
+    def __rmul__(self, other):
+        return "reflected"
+
+
+assert ctypes.c_int * ReflectedLength() == "reflected"
+
+
 class ArrayField(ctypes.Structure):
     _fields_ = [("prefix", ctypes.c_int), ("values", ctypes.c_int * 3)]
 
@@ -518,6 +554,7 @@ for raw_view in raw_views:
         memoryview(dest).cast("B")[::2] = memoryview(source).cast("B")[:4]
         assert list(raw_owner) == [0, 1, 1, 3, 2, 5, 3, 7]
 
+
 # A snapshotted source's release hook must run after the destination write.
 raw_owner[:] = range(8)
 destination_buffer = bytearray(b"abcdefgh")