Add buffer support to c-api (#8071)

* Add buffer support to c-api

* Review

* Refactor PyBuffer

* Manual implement `PyBuffer_GetPointer`

* Review

* Store in 1 Box

crates/capi/src/lib.rs

@@ -32,6 +32,7 @@ pub mod moduleobject;
 pub mod object;
 pub mod objimpl;
 pub mod osmodule;
+pub mod pybuffer;
 pub mod pycapsule;
 pub mod pyerrors;
 pub mod pyframe;

crates/capi/src/pybuffer.rs

@@ -0,0 +1,334 @@
+use crate::PyObject;
+use crate::pystate::with_vm;
+use crate::util::FfiPtrExt;
+use alloc::borrow::Cow;
+use alloc::ffi::CString;
+use core::ffi::{c_char, c_int, c_void};
+use core::ptr::{self, NonNull};
+use rustpython_vm::protocol::{BufferFlags, PyBuffer};
+
+#[repr(C)]
+#[derive(Default)]
+pub struct Py_buffer {
+    pub buf: *mut c_void,
+    pub obj: *mut PyObject,
+    pub len: isize,
+    pub itemsize: isize,
+    pub readonly: c_int,
+    pub ndim: c_int,
+    pub format: *mut c_char,
+    pub shape: *mut isize,
+    pub strides: *mut isize,
+    pub suboffsets: *mut isize,
+    pub internal: *mut c_void,
+}
+
+impl Py_buffer {
+    fn is_contiguous_for_order(&self, order: u8) -> bool {
+        let desc = unsafe { &(&*self.internal.cast::<BufferInternal>()).buffer.desc };
+        if desc.len == 0 || desc.ndim() <= 1 {
+            return true;
+        }
+
+        if desc.has_suboffsets() {
+            return false;
+        }
+
+        match order {
+            b'C' => desc.is_contiguous(),
+            b'F' => desc.is_fortran_contiguous(),
+            b'A' => desc.is_contiguous() || desc.is_fortran_contiguous(),
+            _ => false,
+        }
+    }
+}
+
+#[allow(dead_code)]
+struct BufferInternal {
+    buffer: PyBuffer,
+    format: Option<CString>,
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyObject_GetBuffer(
+    obj: *mut PyObject,
+    view: *mut Py_buffer,
+    flags: c_int,
+) -> c_int {
+    with_vm(|vm| {
+        let obj = unsafe { obj.assume_borrowed() };
+        let flags = BufferFlags::from_bits_retain(flags as u32);
+        let buffer = PyBuffer::from_object(vm, obj, flags)?;
+
+        if flags.contains(BufferFlags::WRITABLE) && buffer.desc.readonly {
+            return Err(vm.new_buffer_error("Object is not writable"));
+        }
+
+        let format = flags
+            .contains(BufferFlags::FORMAT)
+            .then(|| {
+                CString::new(&*buffer.desc.format)
+                    .map_err(|_| vm.new_system_error("buffer format contains NUL"))
+            })
+            .transpose()?;
+
+        let shape = if flags.contains(BufferFlags::ND) {
+            buffer.desc.shape().as_ptr().cast_mut()
+        } else {
+            ptr::null_mut()
+        };
+
+        let strides = if flags.contains(BufferFlags::STRIDES) {
+            buffer.desc.strides().as_ptr().cast_mut()
+        } else {
+            ptr::null_mut()
+        };
+
+        let suboffsets = if flags.contains(BufferFlags::INDIRECT) {
+            buffer.desc.suboffsets().as_ptr().cast_mut()
+        } else {
+            ptr::null_mut()
+        };
+
+        let mut buffer_view = Py_buffer {
+            len: buffer.desc.len as _,
+            itemsize: buffer.desc.itemsize as _,
+            readonly: buffer.desc.readonly.into(),
+            ndim: buffer.desc.ndim() as _,
+            format: format.as_ref().map_or_default(|s| s.as_ptr().cast_mut()),
+            shape,
+            strides,
+            suboffsets,
+            ..Default::default()
+        };
+
+        let c_contig = buffer.desc.is_contiguous();
+        let f_contig = buffer.desc.is_fortran_contiguous();
+
+        if flags.contains(BufferFlags::C_CONTIGUOUS) && !c_contig {
+            return Err(vm.new_buffer_error("Object is not C-contiguous"));
+        }
+        if flags.contains(BufferFlags::F_CONTIGUOUS) && !f_contig {
+            return Err(vm.new_buffer_error("Object is not Fortran-contiguous"));
+        }
+        if flags.contains(BufferFlags::ANY_CONTIGUOUS) && !(c_contig || f_contig) {
+            return Err(vm.new_buffer_error("Object is not contiguous"));
+        }
+
+        buffer_view.obj = obj.as_raw().cast_mut();
+        buffer_view.buf = unsafe {
+            buffer
+                .obj_bytes()
+                .as_ptr()
+                .cast_mut()
+                .offset(buffer.desc.offset)
+        }
+        .cast();
+        buffer_view.internal = Box::into_raw(Box::new(BufferInternal { buffer, format })).cast();
+
+        unsafe {
+            ptr::replace(view, buffer_view);
+        }
+        Ok(())
+    })
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyBuffer_Release(view: *mut Py_buffer) {
+    let view = unsafe { &mut *view };
+    let internal = unsafe {
+        NonNull::new(view.internal.cast::<BufferInternal>()).map(|ptr| Box::from_non_null(ptr))
+    };
+    drop(internal);
+    core::mem::take(view);
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyBuffer_IsContiguous(view: *const Py_buffer, fort: c_char) -> c_int {
+    unsafe { &*view }
+        .is_contiguous_for_order((fort as u8).to_ascii_uppercase())
+        .into()
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyBuffer_GetPointer(
+    view: *const Py_buffer,
+    indices: *const isize,
+) -> *mut c_void {
+    with_vm(|vm| {
+        let view = unsafe { &*view };
+
+        if view.strides.is_null() && view.shape.is_null() {
+            let i0 = unsafe { *indices };
+            let delta = i0
+                .checked_mul(view.itemsize)
+                .ok_or_else(|| vm.new_system_error("index * itemsize overflow"))?;
+            return unsafe { Ok(view.buf.cast::<u8>().offset(delta).cast()) };
+        }
+
+        let ndim: usize = view
+            .ndim
+            .try_into()
+            .map_err(|_| vm.new_system_error("view.ndim does not fit usize"))?;
+
+        let idx = unsafe { core::slice::from_raw_parts(indices, ndim) };
+
+        let strides: Cow<'_, _> = unsafe { view.strides.as_ref() }.map_or_else(
+            || {
+                let shape = unsafe { core::slice::from_raw_parts(view.shape, ndim) };
+                let mut strides = vec![0; shape.len()];
+                let mut stride = view.itemsize;
+                for (ix, dim) in shape.iter().copied().enumerate().rev() {
+                    strides[ix] = stride;
+                    stride = stride
+                        .checked_mul(dim)
+                        .ok_or_else(|| vm.new_system_error("stride overflow"))?;
+                }
+                Ok(strides.into())
+            },
+            |strides| unsafe { Ok(core::slice::from_raw_parts(strides, ndim).into()) },
+        )?;
+
+        let suboffsets = unsafe {
+            view.suboffsets
+                .as_ref()
+                .map(|suboffsets| core::slice::from_raw_parts(suboffsets, ndim))
+        };
+
+        let mut ptr_u8 = view.buf.cast::<u8>();
+        for (dim, index) in idx.iter().enumerate() {
+            let delta = index
+                .checked_mul(strides[dim])
+                .ok_or_else(|| vm.new_system_error("index * stride overflow"))?;
+            ptr_u8 = unsafe { ptr_u8.offset(delta) };
+            if let Some(suboffsets) = suboffsets {
+                let suboffset = suboffsets[dim];
+                if suboffset >= 0 {
+                    let inner = unsafe { core::ptr::read_unaligned(ptr_u8.cast::<*mut u8>()) };
+                    ptr_u8 = unsafe { inner.offset(suboffset) };
+                }
+            }
+        }
+
+        Ok(ptr_u8.cast())
+    })
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyBuffer_ToContiguous(
+    buf: *mut c_void,
+    view: *const Py_buffer,
+    len: isize,
+    order: c_char,
+) -> c_int {
+    with_vm(|vm| {
+        let view = unsafe { &*view };
+
+        if !view.is_contiguous_for_order((order as u8).to_ascii_uppercase()) {
+            return Err(vm.new_buffer_error(
+                "PyBuffer_ToContiguous only supports contiguous exported buffers",
+            ));
+        }
+
+        if len != view.len {
+            return Err(vm.new_buffer_error("len must match view->len"));
+        }
+
+        let len: usize = len
+            .try_into()
+            .map_err(|_| vm.new_system_error("buffer len does not fit usize"))?;
+
+        let src = unsafe { core::slice::from_raw_parts(view.buf.cast::<u8>(), len) };
+        let dst = unsafe { core::slice::from_raw_parts_mut(buf.cast::<u8>(), len) };
+
+        dst.copy_from_slice(src);
+        Ok(())
+    })
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyBuffer_FromContiguous(
+    view: *const Py_buffer,
+    buf: *const c_void,
+    len: isize,
+    order: c_char,
+) -> c_int {
+    with_vm(|vm| {
+        let view = unsafe { &*view };
+
+        if view.readonly != 0 {
+            return Err(vm.new_buffer_error("cannot write into readonly buffer"));
+        }
+
+        if !view.is_contiguous_for_order((order as u8).to_ascii_uppercase()) {
+            return Err(vm.new_buffer_error(
+                "PyBuffer_FromContiguous only supports contiguous exported buffers",
+            ));
+        }
+
+        let len: usize = len
+            .min(view.len)
+            .try_into()
+            .map_err(|_| vm.new_system_error("len does not fit usize"))?;
+
+        let src = unsafe { core::slice::from_raw_parts(buf.cast::<u8>(), len) };
+        let dst = unsafe { core::slice::from_raw_parts_mut(view.buf.cast::<u8>(), len) };
+        dst.copy_from_slice(src);
+        Ok(())
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use pyo3::buffer::PyBuffer;
+    use pyo3::prelude::*;
+    use pyo3::types::{PyByteArray, PyBytes};
+
+    #[test]
+    fn object_getbuffer_basic_and_release() {
+        Python::attach(|py| {
+            let bytes = PyBytes::new(py, b"hello");
+            let buffer = PyBuffer::<u8>::get(&bytes).unwrap();
+            assert_eq!(buffer.dimensions(), 1);
+            assert_eq!(buffer.item_count(), 5);
+            assert_eq!(buffer.to_vec(py).unwrap(), b"hello");
+        });
+    }
+
+    #[test]
+    fn contiguous_copy_roundtrip() {
+        Python::attach(|py| {
+            let src = PyBytes::new(py, b"abcde");
+            let buffer = PyBuffer::<u8>::get(&src).unwrap();
+            let mut out = [0u8; 5];
+            buffer.copy_to_slice(py, &mut out).unwrap();
+            assert_eq!(&out, b"abcde");
+        });
+    }
+
+    #[test]
+    fn is_contiguous_and_get_pointer() {
+        Python::attach(|py| {
+            let bytes = PyBytes::new(py, b"xyz");
+            let buffer = PyBuffer::<u8>::get(&bytes).unwrap();
+            assert!(buffer.is_c_contiguous());
+
+            let p = buffer.get_ptr(&[1]);
+            assert!(!p.is_null());
+            unsafe { assert_eq!(*(p.cast::<u8>()), b'y') };
+        });
+    }
+
+    #[test]
+    fn writable_bytearray() {
+        Python::attach(|py| {
+            let bytearray = PyByteArray::new(py, b"hello");
+            let buffer = PyBuffer::<u8>::get(&bytearray).unwrap();
+            assert_eq!(buffer.dimensions(), 1);
+            assert_eq!(buffer.item_count(), 5);
+            buffer.as_mut_slice(py).unwrap()[0].replace(b'H');
+            drop(buffer);
+            assert_eq!(bytearray.to_vec(), b"Hello");
+        });
+    }
+}

crates/vm/src/builtins/memory.rs

@@ -301,7 +301,7 @@ impl PyMemoryView {
                 vm.new_not_implemented_error("multi-dimensional sub-views are not implemented")
             );
         }
-        let (shape, _, _) = self.desc.dim_desc[0];
+        let shape = self.desc.shape()[0] as usize;
         // ptr_from_index
         let index = i
             .wrapped_at(shape)
@@ -327,7 +327,7 @@ impl PyMemoryView {
         if self.desc.ndim() != 1 {
             return Err(vm.new_not_implemented_error("sub-views are not implemented"));
         }
-        let (shape, _, _) = self.desc.dim_desc[0];
+        let shape = self.desc.shape()[0] as usize;
         // ptr_from_index
         let index = i
             .wrapped_at(shape)
@@ -412,7 +412,7 @@ impl PyMemoryView {
     }
 
     fn init_len(&mut self) {
-        let product: usize = self.desc.dim_desc.iter().map(|x| x.0).product();
+        let product: usize = self.desc.shape().iter().map(|&x| x as usize).product();
         self.desc.len = product * self.desc.itemsize;
     }
 
@@ -420,34 +420,36 @@ impl PyMemoryView {
     /// outside `dim` is reached through a pointer, in which case its suboffset
     /// does.
     fn adjust_position(&mut self, dim: usize, delta: isize) {
-        match self.desc.dim_desc[..dim]
+        match self.desc.suboffsets()[..dim]
             .iter()
-            .rposition(|&(_, _, suboffset)| suboffset != 0)
+            .rposition(|&suboffset| suboffset != -1)
         {
-            Some(n) => self.desc.dim_desc[n].2 += delta,
+            Some(n) => self.desc.suboffsets_mut()[n] += delta,
             None => self.desc.offset += delta,
         }
     }
 
     fn init_range(&mut self, range: Range<usize>, dim: usize) {
-        let (shape, stride, _) = self.desc.dim_desc[dim];
+        let shape = self.desc.shape()[dim] as usize;
+        let stride = self.desc.strides()[dim];
         debug_assert!(shape >= range.len());
 
         self.adjust_position(dim, stride * range.start as isize);
-        self.desc.dim_desc[dim].0 = range.len();
+        self.desc.shape_mut()[dim] = range.len() as isize;
     }
 
     // init_slice
     fn init_slice(&mut self, slice: &Py<PySlice>, dim: usize, vm: &VirtualMachine) -> PyResult<()> {
-        let (shape, stride, _) = self.desc.dim_desc[dim];
+        let shape = self.desc.shape()[dim] as usize;
+        let stride = self.desc.strides()[dim];
         let slice = slice.to_saturated(vm)?;
         let (start, slice_len) = slice.adjust_indices_start(shape);
 
         // Repeated slicing multiplies the stride by the step every time, which
         // overflows after about twenty rounds; C wraps there and so does this.
         self.adjust_position(dim, stride.wrapping_mul(start));
-        self.desc.dim_desc[dim].0 = slice_len;
-        self.desc.dim_desc[dim].1 = stride.wrapping_mul(slice.step());
+        self.desc.shape_mut()[dim] = slice_len as isize;
+        self.desc.strides_mut()[dim] = stride.wrapping_mul(slice.step());
 
         Ok(())
     }
@@ -459,7 +461,13 @@ impl PyMemoryView {
         dim: usize,
         vm: &VirtualMachine,
     ) -> PyResult<PyListRef> {
-        let (shape, stride, suboffset) = self.desc.dim_desc[dim];
+        let shape = self.desc.shape()[dim] as usize;
+        let stride = self.desc.strides()[dim];
+        let suboffset = if self.desc.suboffsets()[dim] == -1 {
+            0
+        } else {
+            self.desc.suboffsets()[dim]
+        };
         if dim + 1 == self.desc.ndim() {
             let mut v = Vec::with_capacity(shape);
             for _ in 0..shape {
@@ -731,7 +739,7 @@ impl PyMemoryView {
             Err(vm.new_type_error("0-dim memory has no length"))
         } else {
             // shape for dim[0]
-            Ok(self.desc.dim_desc[0].0)
+            Ok(self.desc.shape()[0] as usize)
         }
     }
 
@@ -760,14 +768,14 @@ impl PyMemoryView {
             released: AtomicCell::new(false),
             restricted: AtomicCell::new(false),
             format_spec,
-            desc: BufferDescriptor {
-                len: self.desc.len,
-                offset: self.desc.offset,
-                readonly: self.desc.readonly,
+            desc: BufferDescriptor::from_dim_desc(
+                self.desc.len,
+                self.desc.offset,
+                self.desc.readonly,
                 itemsize,
-                format: format_str.to_owned().into(),
-                dim_desc: vec![(self.desc.len / itemsize, itemsize as isize, 0)],
-            },
+                format_str.to_owned().into(),
+                vec![(self.desc.len / itemsize, itemsize as isize, 0)],
+            ),
             hash: OnceCell::new(),
             exports: AtomicCell::new(0),
         };
@@ -903,9 +911,9 @@ impl Py<PyMemoryView> {
         self.try_not_released(vm)?;
         Ok(vm.ctx.new_tuple(
             self.desc
-                .dim_desc
+                .shape()
                 .iter()
-                .map(|(shape, _, _)| shape.to_pyobject(vm))
+                .map(|shape| (*shape as usize).to_pyobject(vm))
                 .collect(),
         ))
     }
@@ -915,9 +923,9 @@ impl Py<PyMemoryView> {
         self.try_not_released(vm)?;
         Ok(vm.ctx.new_tuple(
             self.desc
-                .dim_desc
+                .strides()
                 .iter()
-                .map(|(_, stride, _)| stride.to_pyobject(vm))
+                .map(|stride| stride.to_pyobject(vm))
                 .collect(),
         ))
     }
@@ -927,15 +935,15 @@ impl Py<PyMemoryView> {
         self.try_not_released(vm)?;
         let has_suboffsets = self
             .desc
-            .dim_desc
+            .suboffsets()
             .iter()
-            .any(|(_, _, suboffset)| *suboffset != 0);
+            .any(|suboffset| *suboffset != -1);
         if has_suboffsets {
             Ok(vm.ctx.new_tuple(
                 self.desc
-                    .dim_desc
+                    .suboffsets()
                     .iter()
-                    .map(|(_, _, suboffset)| suboffset.to_pyobject(vm))
+                    .map(|suboffset| suboffset.to_pyobject(vm))
                     .collect(),
             ))
         } else {
@@ -1055,7 +1063,7 @@ impl Py<PyMemoryView> {
                 vm.new_not_implemented_error("multi-dimensional sub-views are not implemented")
             );
         }
-        let len = self.desc.dim_desc[0].0;
+        let len = self.desc.shape()[0] as usize;
         let mut count = 0;
         for i in 0..len {
             let item = self.getitem_by_idx(i as isize, vm)?;
@@ -1074,7 +1082,7 @@ impl Py<PyMemoryView> {
                 vm.new_not_implemented_error("multi-dimensional sub-views are not implemented")
             );
         }
-        let len = self.desc.dim_desc[0].0;
+        let len = self.desc.shape()[0] as usize;
         let MemoryIndexArgs { value, start, stop } = args;
 
         let start = if start < 0 {
@@ -1148,12 +1156,14 @@ impl Py<PyMemoryView> {
                         vm.new_type_error("memoryview: product(shape) * itemsize != buffer size")
                     );
                 }
-                other.desc.dim_desc = vec![];
+                other
+                    .desc
+                    .set_dimensions(Box::new([]), Box::new([]), Box::new([]));
                 return Ok(other.into_ref(&vm.ctx));
             }
 
             let mut product_shape = itemsize;
-            let mut dim_descriptor = Vec::with_capacity(shape_ndim);
+            let mut shape_descriptor = Vec::with_capacity(shape_ndim);
 
             for x in shape {
                 let x = x
@@ -1174,12 +1184,13 @@ impl Py<PyMemoryView> {
                     return Err(vm.new_value_error("memoryview.cast(): product(shape) > SSIZE_MAX"));
                 }
                 product_shape *= x;
-                dim_descriptor.push((x, 0, 0));
+                shape_descriptor.push(x as isize);
             }
 
-            dim_descriptor.last_mut().unwrap().1 = itemsize as isize;
-            for i in (0..dim_descriptor.len() - 1).rev() {
-                dim_descriptor[i].1 = dim_descriptor[i + 1].1 * dim_descriptor[i + 1].0 as isize;
+            let mut strides = vec![0isize; shape_descriptor.len()];
+            strides[shape_descriptor.len() - 1] = itemsize as isize;
+            for i in (0..strides.len() - 1).rev() {
+                strides[i] = strides[i + 1] * shape_descriptor[i + 1];
             }
 
             if product_shape != other.desc.len {
@@ -1188,7 +1199,11 @@ impl Py<PyMemoryView> {
                 );
             }
 
-            other.desc.dim_desc = dim_descriptor;
+            other.desc.set_dimensions(
+                shape_descriptor.into_boxed_slice(),
+                strides.into_boxed_slice(),
+                vec![-1; shape_ndim].into_boxed_slice(),
+            );
 
             Ok(other.into_ref(&vm.ctx))
         } else {
@@ -1670,8 +1685,8 @@ fn is_equiv_shape(a: &BufferDescriptor, b: &BufferDescriptor) -> bool {
         return false;
     }
 
-    let a_iter = a.dim_desc.iter().map(|x| x.0);
-    let b_iter = b.dim_desc.iter().map(|x| x.0);
+    let a_iter = a.shape().iter().copied();
+    let b_iter = b.shape().iter().copied();
     for (a_shape, b_shape) in a_iter.zip(b_iter) {
         if a_shape != b_shape {
             return false;

crates/vm/src/protocol/buffer.rs

@@ -403,25 +403,110 @@ pub struct BufferDescriptor {
     pub readonly: bool,
     pub itemsize: usize,
     pub format: Cow<'static, str>,
-    /// (shape, stride, suboffset) for each dimension. A non-zero suboffset means
-    /// the dimension is reached through a pointer; slicing never introduces one.
-    pub dim_desc: Vec<(usize, isize, isize)>,
+    /// Number of dimensions.
+    ndim: usize,
+    /// Shape, strides and suboffsets packed as 3 contiguous slices of length `ndim`.
+    dims: Box<[isize]>,
     // TODO: flags
 }
 
 impl BufferDescriptor {
     #[must_use]
-    pub fn simple(bytes_len: usize, readonly: bool) -> Self {
+    pub fn from_dim_desc(
+        len: usize,
+        offset: isize,
+        readonly: bool,
+        itemsize: usize,
+        format: Cow<'static, str>,
+        dim_desc: Vec<(usize, isize, isize)>,
+    ) -> Self {
+        let ndim = dim_desc.len();
+        let mut dims = Vec::with_capacity(ndim * 3);
+        dims.extend(dim_desc.iter().map(|(shape, _, _)| *shape as isize));
+        dims.extend(dim_desc.iter().map(|(_, stride, _)| *stride));
+        dims.extend(
+            dim_desc
+                .iter()
+                .map(|(_, _, suboffset)| if *suboffset == 0 { -1 } else { *suboffset }),
+        );
         Self {
-            len: bytes_len,
-            offset: 0,
+            len,
+            offset,
             readonly,
-            itemsize: 1,
-            format: Cow::Borrowed("B"),
-            dim_desc: vec![(bytes_len, 1, 0)],
+            itemsize,
+            format,
+            ndim,
+            dims: dims.into_boxed_slice(),
         }
     }
 
+    #[must_use]
+    pub fn shape(&self) -> &[isize] {
+        debug_assert_eq!(self.dims.len(), self.ndim * 3);
+        &self.dims[..self.ndim]
+    }
+
+    #[must_use]
+    pub fn strides(&self) -> &[isize] {
+        debug_assert_eq!(self.dims.len(), self.ndim * 3);
+        &self.dims[self.ndim..self.ndim * 2]
+    }
+
+    #[must_use]
+    pub fn suboffsets(&self) -> &[isize] {
+        debug_assert_eq!(self.dims.len(), self.ndim * 3);
+        &self.dims[self.ndim * 2..]
+    }
+
+    pub fn shape_mut(&mut self) -> &mut [isize] {
+        debug_assert_eq!(self.dims.len(), self.ndim * 3);
+        &mut self.dims[..self.ndim]
+    }
+
+    pub fn strides_mut(&mut self) -> &mut [isize] {
+        debug_assert_eq!(self.dims.len(), self.ndim * 3);
+        &mut self.dims[self.ndim..self.ndim * 2]
+    }
+
+    pub fn suboffsets_mut(&mut self) -> &mut [isize] {
+        debug_assert_eq!(self.dims.len(), self.ndim * 3);
+        &mut self.dims[self.ndim * 2..]
+    }
+
+    pub fn set_dimensions(
+        &mut self,
+        shape: Box<[isize]>,
+        strides: Box<[isize]>,
+        suboffsets: Box<[isize]>,
+    ) {
+        let ndim = shape.len();
+        debug_assert_eq!(strides.len(), ndim);
+        debug_assert_eq!(suboffsets.len(), ndim);
+
+        let shape = Vec::from(shape);
+        let strides = Vec::from(strides);
+        let suboffsets = Vec::from(suboffsets);
+        let mut dims = Vec::with_capacity(ndim * 3);
+        dims.extend(shape);
+        dims.extend(strides);
+        dims.extend(suboffsets);
+
+        self.ndim = ndim;
+        self.dims = dims.into_boxed_slice();
+    }
+
+    #[must_use]
+    pub fn simple(bytes_len: usize, readonly: bool) -> Self {
+        Self::from_dim_desc(
+            bytes_len,
+            0,
+            readonly,
+            1,
+            Cow::Borrowed("B"),
+            vec![(bytes_len, 1, 0)],
+        )
+    }
+
     #[must_use]
     pub fn format(
         bytes_len: usize,
@@ -429,14 +514,14 @@ impl BufferDescriptor {
         itemsize: usize,
         format: Cow<'static, str>,
     ) -> Self {
-        Self {
-            len: bytes_len,
-            offset: 0,
+        Self::from_dim_desc(
+            bytes_len,
+            0,
             readonly,
             itemsize,
             format,
-            dim_desc: vec![(bytes_len / itemsize, itemsize as isize, 0)],
-        }
+            vec![(bytes_len / itemsize, itemsize as isize, 0)],
+        )
     }
 
     /// The descriptor an exporter hands to a consumer that asked for `flags`:
@@ -460,14 +545,19 @@ impl BufferDescriptor {
             // A request this flat is refused unless the layout is C-contiguous, so
             // one dimension addresses the same bytes.
             let shape = desc.len.checked_div(desc.itemsize).unwrap_or(0);
-            desc.dim_desc = vec![(shape, desc.itemsize as isize, 0)];
+            desc.set_dimensions(
+                Box::new([shape as isize]),
+                Box::new([desc.itemsize as isize]),
+                Box::new([-1]),
+            );
         } else if !flags.contains(BufferFlags::STRIDES) {
             // Shape survives but strides do not, which means C order.
             let mut stride = desc.itemsize as isize;
-            for (shape, dim_stride, suboffset) in desc.dim_desc.iter_mut().rev() {
-                *dim_stride = stride;
-                *suboffset = 0;
-                stride *= *shape as isize;
+            for i in (0..desc.ndim()).rev() {
+                let dim_shape = desc.shape()[i];
+                desc.strides_mut()[i] = stride;
+                desc.suboffsets_mut()[i] = -1;
+                stride *= dim_shape;
             }
         }
         desc
@@ -490,10 +580,9 @@ impl BufferDescriptor {
             debug_assert_eq!(self.itemsize, self.len);
         } else {
             let mut shape_product = 1;
-            let has_zero_dim = self.dim_desc.iter().any(|(s, _, _)| *s == 0);
-            for (shape, stride, suboffset) in self.dim_desc.iter().copied() {
-                shape_product *= shape;
-                debug_assert!(suboffset >= 0);
+            let has_zero_dim = self.shape().contains(&0);
+            for (&shape, &stride) in self.shape().iter().zip(self.strides().iter()) {
+                shape_product *= shape as usize;
                 // For empty arrays (any dimension is 0), strides can be 0
                 if !has_zero_dim {
                     debug_assert_ne!(stride, 0);
@@ -506,7 +595,8 @@ impl BufferDescriptor {
 
     #[must_use]
     pub fn ndim(&self) -> usize {
-        self.dim_desc.len()
+        debug_assert_eq!(self.dims.len(), self.ndim * 3);
+        self.ndim
     }
 
     /// Whether the elements are laid out in row-major order. _IsCContiguous
@@ -516,11 +606,11 @@ impl BufferDescriptor {
             return true;
         }
         let mut sd = self.itemsize;
-        for (shape, stride, _) in self.dim_desc.iter().copied().rev() {
+        for (&shape, &stride) in self.shape().iter().zip(self.strides().iter()).rev() {
             if shape > 1 && stride != sd as isize {
                 return false;
             }
-            sd *= shape;
+            sd *= shape as usize;
         }
         true
     }
@@ -534,11 +624,11 @@ impl BufferDescriptor {
             return true;
         }
         let mut sd = self.itemsize;
-        for (shape, stride, _) in self.dim_desc.iter().copied() {
+        for (&shape, &stride) in self.shape().iter().zip(self.strides().iter()) {
             if shape > 1 && stride != sd as isize {
                 return false;
             }
-            sd *= shape;
+            sd *= shape as usize;
         }
         true
     }
@@ -561,23 +651,27 @@ impl BufferDescriptor {
     /// The same shape, format and item size, laid out in C order from byte zero.
     #[must_use]
     pub fn contiguous(&self) -> Self {
-        let itemsize = self.itemsize;
-        let mut dim_desc = self.dim_desc.clone();
-        if let Some((_, stride, suboffset)) = dim_desc.last_mut() {
-            *stride = itemsize as isize;
-            *suboffset = 0;
-        }
-        for i in (1..dim_desc.len()).rev() {
-            dim_desc[i - 1].1 = dim_desc[i].1 * dim_desc[i].0 as isize;
-            dim_desc[i - 1].2 = 0;
+        let mut strides = vec![0isize; self.ndim()];
+        let mut stride = self.itemsize as isize;
+        for i in (0..self.ndim()).rev() {
+            strides[i] = stride;
+            stride *= self.shape()[i];
         }
+        let shape = self.shape().to_vec().into_boxed_slice();
         Self {
             len: self.len,
             offset: 0,
             readonly: self.readonly,
             itemsize: self.itemsize,
             format: self.format.clone(),
-            dim_desc,
+            ndim: self.ndim,
+            dims: {
+                let mut dims = Vec::with_capacity(self.ndim * 3);
+                dims.extend(shape.iter().copied());
+                dims.extend(strides.iter().copied());
+                dims.extend(vec![-1; self.ndim]);
+                dims.into_boxed_slice()
+            },
         }
     }
 
@@ -585,9 +679,7 @@ impl BufferDescriptor {
     /// stepping, the layout `PyBUF_INDIRECT` describes.
     #[must_use]
     pub fn has_suboffsets(&self) -> bool {
-        self.dim_desc
-            .iter()
-            .any(|(_, _, suboffset)| *suboffset != 0)
+        self.suboffsets().iter().any(|&suboffset| suboffset != -1)
     }
 
     /// this function do not check the bound
@@ -595,12 +687,13 @@ impl BufferDescriptor {
     #[must_use]
     pub fn fast_position(&self, indices: &[usize]) -> isize {
         let mut pos = self.offset;
-        for (i, (_, stride, suboffset)) in indices
-            .iter()
-            .copied()
-            .zip_eq(self.dim_desc.iter().copied())
-        {
-            pos += i as isize * stride + suboffset;
+        for (dim, i) in indices.iter().copied().enumerate() {
+            let suboffset = if self.suboffsets()[dim] == -1 {
+                0
+            } else {
+                self.suboffsets()[dim]
+            };
+            pos += i as isize * self.strides()[dim] + suboffset;
         }
         pos
     }
@@ -611,7 +704,19 @@ impl BufferDescriptor {
         for (dim, (i, (shape, stride, suboffset))) in indices
             .iter()
             .copied()
-            .zip_eq(self.dim_desc.iter().copied())
+            .zip_eq(
+                self.shape()
+                    .iter()
+                    .zip(self.strides().iter())
+                    .zip(self.suboffsets().iter())
+                    .map(|((&shape, &stride), &suboffset)| {
+                        (
+                            shape as usize,
+                            stride,
+                            if suboffset == -1 { 0 } else { suboffset },
+                        )
+                    }),
+            )
             .enumerate()
         {
             // The dimension is named the way a person counts it. lookup_dimension
@@ -664,15 +769,17 @@ impl BufferDescriptor {
             let pos = self.offset
                 + indices
                     .iter()
-                    .zip_eq(self.dim_desc.iter())
-                    .map(|(&i, &(_, stride, suboffset))| i as isize * stride + suboffset)
+                    .zip_eq(self.strides().iter().zip(self.suboffsets().iter()))
+                    .map(|(&i, (&stride, &suboffset))| {
+                        i as isize * stride + if suboffset == -1 { 0 } else { suboffset }
+                    })
                     .sum::<isize>();
             f(pos..pos + self.itemsize as isize);
 
             let mut dim = 0;
             loop {
                 indices[dim] += 1;
-                if indices[dim] < self.dim_desc[dim].0 {
+                if (indices[dim] as isize) < self.shape()[dim] {
                     break;
                 }
                 indices[dim] = 0;
@@ -688,7 +795,13 @@ impl BufferDescriptor {
     where
         F: FnMut(Range<isize>),
     {
-        let (shape, stride, suboffset) = self.dim_desc[dim];
+        let shape = self.shape()[dim] as usize;
+        let stride = self.strides()[dim];
+        let suboffset = if self.suboffsets()[dim] == -1 {
+            0
+        } else {
+            self.suboffsets()[dim]
+        };
         if dim + 1 == self.ndim() {
             if CONTIGUOUS {
                 f(index..index + (shape * self.itemsize) as isize);
@@ -743,9 +856,21 @@ impl BufferDescriptor {
     ) where
         F: FnMut(Range<isize>, Range<isize>) -> bool,
     {
-        let (shape, a_stride, a_suboffset) = self.dim_desc[dim];
-        let (_b_shape, b_stride, b_suboffset) = other.dim_desc[dim];
-        debug_assert_eq!(shape, _b_shape);
+        let shape = self.shape()[dim] as usize;
+        let a_stride = self.strides()[dim];
+        let a_suboffset = if self.suboffsets()[dim] == -1 {
+            0
+        } else {
+            self.suboffsets()[dim]
+        };
+        let b_shape = other.shape()[dim] as usize;
+        let b_stride = other.strides()[dim];
+        let b_suboffset = if other.suboffsets()[dim] == -1 {
+            0
+        } else {
+            other.suboffsets()[dim]
+        };
+        debug_assert_eq!(shape, b_shape);
         if dim + 1 == self.ndim() {
             if CONTIGUOUS {
                 if f(
@@ -786,13 +911,13 @@ impl BufferDescriptor {
 
     #[must_use]
     fn is_last_dim_contiguous(&self) -> bool {
-        let (_, stride, suboffset) = self.dim_desc[self.ndim() - 1];
-        suboffset == 0 && stride == self.itemsize as isize
+        let i = self.ndim() - 1;
+        self.suboffsets()[i] == -1 && self.strides()[i] == self.itemsize as isize
     }
 
     #[must_use]
     pub fn is_zero_in_shape(&self) -> bool {
-        self.dim_desc.iter().any(|(shape, _, _)| *shape == 0)
+        self.shape().contains(&0)
     }
 
     // TODO: support column-major order

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

@@ -1060,14 +1060,14 @@ impl AsBuffer for PyCArray {
             }
             dim_desc.reverse();
 
-            BufferDescriptor {
-                offset: 0,
-                len: buffer_len,
-                readonly: false,
+            BufferDescriptor::from_dim_desc(
+                buffer_len,
+                0,
+                false,
                 itemsize,
-                format: alloc::borrow::Cow::Owned(fmt),
+                alloc::borrow::Cow::Owned(fmt),
                 dim_desc,
-            }
+            )
         } else {
             // Fallback to simple buffer if no format/shape info
             BufferDescriptor::simple(buffer_len, false)

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

@@ -1663,14 +1663,7 @@ impl AsBuffer for PyCFuncPtr {
         } else {
             (Cow::Borrowed(pointer_format()), pointer_size())
         };
-        let desc = BufferDescriptor {
-            offset: 0,
-            len: itemsize,
-            readonly: false,
-            itemsize,
-            format,
-            dim_desc: vec![],
-        };
+        let desc = BufferDescriptor::from_dim_desc(itemsize, 0, false, itemsize, format, vec![]);
         let buf = PyBuffer::new(zelf.to_owned().into(), desc, &CDATA_BUFFER_METHODS);
         Ok(buf)
     }

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

@@ -781,14 +781,7 @@ impl AsBuffer for PyCPointer {
             .map_or(Cow::Borrowed("&B"), Cow::Owned);
         let itemsize = stg_info.size;
         // Pointer types are scalars with ndim=0, shape=()
-        let desc = BufferDescriptor {
-            offset: 0,
-            len: itemsize,
-            readonly: false,
-            itemsize,
-            format,
-            dim_desc: vec![],
-        };
+        let desc = BufferDescriptor::from_dim_desc(itemsize, 0, false, itemsize, format, vec![]);
         let buf = PyBuffer::new(zelf.to_owned().into(), desc, &CDATA_BUFFER_METHODS);
         Ok(buf)
     }

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

@@ -1288,14 +1288,7 @@ impl AsBuffer for PyCSimple {
             .map_or(Cow::Borrowed("B"), Cow::Owned);
         let itemsize = stg_info.size;
         // Simple types are scalars with ndim=0, shape=()
-        let desc = BufferDescriptor {
-            offset: 0,
-            len: itemsize,
-            readonly: false,
-            itemsize,
-            format,
-            dim_desc: vec![],
-        };
+        let desc = BufferDescriptor::from_dim_desc(itemsize, 0, false, itemsize, format, vec![]);
         let buf = PyBuffer::new(zelf.to_owned().into(), desc, &CDATA_BUFFER_METHODS);
         Ok(buf)
     }

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

@@ -823,14 +823,14 @@ impl AsBuffer for PyCStructure {
         // Structure: ndim=0, shape=(), itemsize=struct_size
         let buf = PyBuffer::new(
             zelf.to_owned().into(),
-            BufferDescriptor {
-                offset: 0,
-                len: buffer_len,
-                readonly: false,
-                itemsize: buffer_len,
-                format: Cow::Owned(format),
-                dim_desc: vec![], // ndim=0 means empty dim_desc
-            },
+            BufferDescriptor::from_dim_desc(
+                buffer_len,
+                0,
+                false,
+                buffer_len,
+                Cow::Owned(format),
+                vec![], // ndim=0 means empty dim_desc
+            ),
             &CDATA_BUFFER_METHODS,
         );
         Ok(buf)

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

@@ -684,14 +684,14 @@ impl AsBuffer for PyCUnion {
         // Union: ndim=0, shape=(), itemsize=union_size
         let buf = PyBuffer::new(
             zelf.to_owned().into(),
-            BufferDescriptor {
-                offset: 0,
-                len: buffer_len,
-                readonly: false,
-                itemsize: buffer_len,
-                format: Cow::Owned(format),
-                dim_desc: vec![], // ndim=0 means empty dim_desc
-            },
+            BufferDescriptor::from_dim_desc(
+                buffer_len,
+                0,
+                false,
+                buffer_len,
+                Cow::Owned(format),
+                vec![], // ndim=0 means empty dim_desc
+            ),
             &CDATA_BUFFER_METHODS,
         );
         Ok(buf)