From e6cbf3628d5e963034a23799e001bfbdd3ec529b Mon Sep 17 00:00:00 2001 From: Bas Schoenmaeckers <7943856+bschoenmaeckers@users.noreply.github.com> Date: Fri, 9 Oct 2026 13:00:59 +0200 Subject: [PATCH] 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 | 1 + crates/capi/src/pybuffer.rs | 334 ++++++++++++++++++++++ crates/vm/src/builtins/memory.rs | 93 +++--- crates/vm/src/protocol/buffer.rs | 241 ++++++++++++---- crates/vm/src/stdlib/_ctypes/array.rs | 12 +- crates/vm/src/stdlib/_ctypes/function.rs | 9 +- crates/vm/src/stdlib/_ctypes/pointer.rs | 9 +- crates/vm/src/stdlib/_ctypes/simple.rs | 9 +- crates/vm/src/stdlib/_ctypes/structure.rs | 16 +- crates/vm/src/stdlib/_ctypes/union.rs | 16 +- 10 files changed, 597 insertions(+), 143 deletions(-) create mode 100644 crates/capi/src/pybuffer.rs diff --git a/crates/capi/src/lib.rs b/crates/capi/src/lib.rs index 90ef71da7..ea9924200 100644 --- a/crates/capi/src/lib.rs +++ b/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; diff --git a/crates/capi/src/pybuffer.rs b/crates/capi/src/pybuffer.rs new file mode 100644 index 000000000..b20f07c6d --- /dev/null +++ b/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::()).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, +} + +#[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::()).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::().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::(); + 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::(), len) }; + let dst = unsafe { core::slice::from_raw_parts_mut(buf.cast::(), 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::(), len) }; + let dst = unsafe { core::slice::from_raw_parts_mut(view.buf.cast::(), 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::::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::::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::::get(&bytes).unwrap(); + assert!(buffer.is_c_contiguous()); + + let p = buffer.get_ptr(&[1]); + assert!(!p.is_null()); + unsafe { assert_eq!(*(p.cast::()), b'y') }; + }); + } + + #[test] + fn writable_bytearray() { + Python::attach(|py| { + let bytearray = PyByteArray::new(py, b"hello"); + let buffer = PyBuffer::::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"); + }); + } +} diff --git a/crates/vm/src/builtins/memory.rs b/crates/vm/src/builtins/memory.rs index 2c9e8c628..33481ce22 100644 --- a/crates/vm/src/builtins/memory.rs +++ b/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, 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, 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 { - 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 { 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 { 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 { 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 { 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 { 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 { 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 { 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 { ); } - 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; diff --git a/crates/vm/src/protocol/buffer.rs b/crates/vm/src/protocol/buffer.rs index 500ea5992..e8f13e22a 100644 --- a/crates/vm/src/protocol/buffer.rs +++ b/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::(); 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), { - 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, Range) -> 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 diff --git a/crates/vm/src/stdlib/_ctypes/array.rs b/crates/vm/src/stdlib/_ctypes/array.rs index 0504688d8..1ce52e47c 100644 --- a/crates/vm/src/stdlib/_ctypes/array.rs +++ b/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) diff --git a/crates/vm/src/stdlib/_ctypes/function.rs b/crates/vm/src/stdlib/_ctypes/function.rs index 6066310d5..4c9b46a93 100644 --- a/crates/vm/src/stdlib/_ctypes/function.rs +++ b/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) } diff --git a/crates/vm/src/stdlib/_ctypes/pointer.rs b/crates/vm/src/stdlib/_ctypes/pointer.rs index 6a8ebe660..27778a01c 100644 --- a/crates/vm/src/stdlib/_ctypes/pointer.rs +++ b/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) } diff --git a/crates/vm/src/stdlib/_ctypes/simple.rs b/crates/vm/src/stdlib/_ctypes/simple.rs index 1a0e12ba4..18b67fca7 100644 --- a/crates/vm/src/stdlib/_ctypes/simple.rs +++ b/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) } diff --git a/crates/vm/src/stdlib/_ctypes/structure.rs b/crates/vm/src/stdlib/_ctypes/structure.rs index bc4b3d54c..ed4cd2548 100644 --- a/crates/vm/src/stdlib/_ctypes/structure.rs +++ b/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) diff --git a/crates/vm/src/stdlib/_ctypes/union.rs b/crates/vm/src/stdlib/_ctypes/union.rs index 2322d6524..8bd012599 100644 --- a/crates/vm/src/stdlib/_ctypes/union.rs +++ b/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)