diff --git a/Lib/test/test_gzip.py b/Lib/test/test_gzip.py index 5ac37fa452..0b5113c20b 100644 --- a/Lib/test/test_gzip.py +++ b/Lib/test/test_gzip.py @@ -88,8 +88,6 @@ def test_write_read_with_pathlike_file(self): self.assertEqual(d, data1 * 51) self.assertIsInstance(f.name, str) - # TODO: RUSTPYTHON - @unittest.expectedFailure # The following test_write_xy methods test that write accepts # the corresponding bytes-like object type as input # and that the data written equals bytes(xy) in all cases. diff --git a/Lib/test/test_memoryview.py b/Lib/test/test_memoryview.py index 3fe85c5035..68d0bfa0ff 100644 --- a/Lib/test/test_memoryview.py +++ b/Lib/test/test_memoryview.py @@ -71,8 +71,6 @@ def setitem(value): m = None self.assertEqual(sys.getrefcount(b), oldrefcount) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_setitem_writable(self): if not self.rw_type: self.skipTest("no writable type to test") diff --git a/extra_tests/snippets/memoryview.py b/extra_tests/snippets/memoryview.py index 121028d418..ca50ece56c 100644 --- a/extra_tests/snippets/memoryview.py +++ b/extra_tests/snippets/memoryview.py @@ -45,6 +45,8 @@ def test_slice(): assert m[::2][1:-1].tobytes() == b'357' assert m[::2][::2].tobytes() == b'159' assert m[::2][1::2].tobytes() == b'37' + assert m[::-1].tobytes() == b'987654321' + assert m[::-2].tobytes() == b'97531' test_slice() @@ -56,9 +58,9 @@ def test_resizable(): m.release() b.append(6) m2 = memoryview(b) - m4 = memoryview(b) + m4 = memoryview(m2) assert_raises(BufferError, lambda: b.append(5)) - m3 = memoryview(b) + m3 = memoryview(m2) assert_raises(BufferError, lambda: b.append(5)) m2.release() assert_raises(BufferError, lambda: b.append(5)) diff --git a/stdlib/src/array.rs b/stdlib/src/array.rs index 8bda4fd6f9..1ce1c9a187 100644 --- a/stdlib/src/array.rs +++ b/stdlib/src/array.rs @@ -3,7 +3,7 @@ pub(crate) use array::make_module; #[pymodule(name = "array")] mod array { use crate::common::{ - borrow::{BorrowedValue, BorrowedValueMut}, + atomic::{self, AtomicUsize}, lock::{ PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard, @@ -20,7 +20,7 @@ mod array { ArgBytesLike, ArgIntoFloat, ArgIterable, IntoPyObject, IntoPyResult, OptionalArg, }, protocol::{ - BufferInternal, BufferOptions, BufferResizeGuard, PyBuffer, PyIterReturn, + BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn, PyMappingMethods, }, sliceable::{PySliceableSequence, PySliceableSequenceMut, SaturatedSlice, SequenceIndex}, @@ -28,10 +28,9 @@ mod array { AsBuffer, AsMapping, Comparable, Constructor, IterNext, IterNextIterable, Iterable, PyComparisonOp, }, - IdProtocol, PyComparisonValue, PyObject, PyObjectRef, PyObjectView, PyRef, PyResult, - PyValue, TryFromObject, TypeProtocol, VirtualMachine, + IdProtocol, PyComparisonValue, PyObject, PyObjectRef, PyObjectView, PyObjectWrap, PyRef, + PyResult, PyValue, TryFromObject, TypeProtocol, VirtualMachine, }; - use crossbeam_utils::atomic::AtomicCell; use itertools::Itertools; use num_traits::ToPrimitive; use std::cmp::Ordering; @@ -615,7 +614,7 @@ mod array { #[derive(Debug, PyValue)] pub struct PyArray { array: PyRwLock, - exports: AtomicCell, + exports: AtomicUsize, } pub type PyArrayRef = PyRef; @@ -624,7 +623,7 @@ mod array { fn from(array: ArrayContentType) -> Self { PyArray { array: PyRwLock::new(array), - exports: AtomicCell::new(0), + exports: AtomicUsize::new(0), } } } @@ -1220,40 +1219,35 @@ mod array { fn as_buffer(zelf: &PyObjectView, _vm: &VirtualMachine) -> PyResult { let array = zelf.read(); let buf = PyBuffer::new( - zelf.as_object().to_owned(), - PyArrayBufferInternal(zelf.to_owned()), - BufferOptions { - readonly: false, - len: array.len(), - itemsize: array.itemsize(), - format: array.typecode_str().into(), - ..Default::default() - }, + zelf.to_owned().into(), + BufferDescriptor::format( + array.len() * array.itemsize(), + false, + array.itemsize(), + array.typecode_str().into(), + ), + &BUFFER_METHODS, ); Ok(buf) } } - #[derive(Debug)] - struct PyArrayBufferInternal(PyRef); - - impl BufferInternal for PyArrayBufferInternal { - fn obj_bytes(&self) -> BorrowedValue<[u8]> { - self.0.get_bytes().into() - } - - fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> { - self.0.get_bytes_mut().into() - } - - fn release(&self) { - self.0.exports.fetch_sub(1); - } - - fn retain(&self) { - self.0.exports.fetch_add(1); - } - } + static BUFFER_METHODS: BufferMethods = BufferMethods { + obj_bytes: |buffer| buffer.obj_as::().get_bytes().into(), + obj_bytes_mut: |buffer| buffer.obj_as::().get_bytes_mut().into(), + release: |buffer| { + buffer + .obj_as::() + .exports + .fetch_sub(1, atomic::Ordering::Release); + }, + retain: |buffer| { + buffer + .obj_as::() + .exports + .fetch_add(1, atomic::Ordering::Release); + }, + }; impl AsMapping for PyArray { fn as_mapping(_zelf: &PyObjectView, _vm: &VirtualMachine) -> PyMappingMethods { @@ -1290,7 +1284,7 @@ mod array { impl Iterable for PyArray { fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { Ok(PyArrayIter { - position: AtomicCell::new(0), + position: AtomicUsize::new(0), array: zelf, } .into_object(vm)) @@ -1302,7 +1296,7 @@ mod array { fn try_resizable(&'a self, vm: &VirtualMachine) -> PyResult { let w = self.write(); - if self.exports.load() == 0 { + if self.exports.load(atomic::Ordering::SeqCst) == 0 { Ok(w) } else { Err(vm.new_buffer_error( @@ -1316,7 +1310,7 @@ mod array { #[pyclass(name = "array_iterator")] #[derive(Debug, PyValue)] pub struct PyArrayIter { - position: AtomicCell, + position: AtomicUsize, array: PyArrayRef, } @@ -1326,7 +1320,7 @@ mod array { impl IterNextIterable for PyArrayIter {} impl IterNext for PyArrayIter { fn next(zelf: &PyObjectView, vm: &VirtualMachine) -> PyResult { - let pos = zelf.position.fetch_add(1); + let pos = zelf.position.fetch_add(1, atomic::Ordering::SeqCst); let r = if let Some(item) = zelf.array.read().getitem_by_idx(pos, vm)? { PyIterReturn::Return(item) } else { diff --git a/vm/src/builtins/bytearray.rs b/vm/src/builtins/bytearray.rs index 5678d6db66..4dc7efe71c 100644 --- a/vm/src/builtins/bytearray.rs +++ b/vm/src/builtins/bytearray.rs @@ -3,13 +3,6 @@ use super::{ PositionIterInternal, PyBytes, PyBytesRef, PyDictRef, PyIntRef, PyStrRef, PyTuple, PyTupleRef, PyTypeRef, }; -use crate::common::{ - borrow::{BorrowedValue, BorrowedValueMut}, - lock::{ - PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyMutex, PyRwLock, PyRwLockReadGuard, - PyRwLockWriteGuard, - }, -}; use crate::{ anystr::{self, AnyStr}, builtins::PyType, @@ -18,9 +11,17 @@ use crate::{ ByteInnerNewOptions, ByteInnerPaddingOptions, ByteInnerSplitOptions, ByteInnerTranslateOptions, DecodeArgs, PyBytesInner, }, + common::{ + atomic::{AtomicUsize, Ordering}, + lock::{ + PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyMutex, PyRwLock, + PyRwLockReadGuard, PyRwLockWriteGuard, + }, + }, function::{ArgBytesLike, ArgIterable, FuncArgs, IntoPyObject, OptionalArg, OptionalOption}, protocol::{ - BufferInternal, BufferOptions, BufferResizeGuard, PyBuffer, PyIterReturn, PyMappingMethods, + BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn, + PyMappingMethods, }, sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex}, types::{ @@ -29,10 +30,9 @@ use crate::{ }, utils::Either, IdProtocol, PyClassDef, PyClassImpl, PyComparisonValue, PyContext, PyObject, PyObjectRef, - PyRef, PyResult, PyValue, TypeProtocol, VirtualMachine, + PyObjectView, PyObjectWrap, PyRef, PyResult, PyValue, TypeProtocol, VirtualMachine, }; use bstr::ByteSlice; -use crossbeam_utils::atomic::AtomicCell; use std::mem::size_of; /// "bytearray(iterable_of_ints) -> bytearray\n\ @@ -50,7 +50,7 @@ use std::mem::size_of; #[derive(Debug, Default)] pub struct PyByteArray { inner: PyRwLock, - exports: AtomicCell, + exports: AtomicUsize, } pub type PyByteArrayRef = PyRef; @@ -63,7 +63,7 @@ impl PyByteArray { fn from_inner(inner: PyBytesInner) -> Self { PyByteArray { inner: PyRwLock::new(inner), - exports: AtomicCell::new(0), + exports: AtomicUsize::new(0), } } @@ -118,6 +118,12 @@ impl PyByteArray { Ok(()) } + #[cfg(debug_assertions)] + #[pyproperty] + fn exports(&self) -> usize { + self.exports.load(Ordering::Relaxed) + } + #[inline] fn inner(&self) -> PyRwLockReadGuard<'_, PyBytesInner> { self.inner.read() @@ -708,36 +714,35 @@ impl Comparable for PyByteArray { } } -impl AsBuffer for PyByteArray { - fn as_buffer(zelf: &crate::PyObjectView, _vm: &VirtualMachine) -> PyResult { - let buffer = PyBuffer::new( - zelf.as_object().to_owned(), - zelf.to_owned(), - BufferOptions { - readonly: false, - len: zelf.len(), - ..Default::default() - }, - ); - Ok(buffer) - } -} - -impl BufferInternal for PyRef { - fn obj_bytes(&self) -> BorrowedValue<[u8]> { - self.borrow_buf().into() - } - - fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> { - PyRwLockWriteGuard::map(self.inner_mut(), |inner| &mut *inner.elements).into() - } - - fn release(&self) { - self.exports.fetch_sub(1); - } +static BUFFER_METHODS: BufferMethods = BufferMethods { + obj_bytes: |buffer| buffer.obj_as::().borrow_buf().into(), + obj_bytes_mut: |buffer| { + PyMappedRwLockWriteGuard::map(buffer.obj_as::().borrow_buf_mut(), |x| { + x.as_mut_slice() + }) + .into() + }, + release: |buffer| { + buffer + .obj_as::() + .exports + .fetch_sub(1, Ordering::Release); + }, + retain: |buffer| { + buffer + .obj_as::() + .exports + .fetch_add(1, Ordering::Release); + }, +}; - fn retain(&self) { - self.exports.fetch_add(1); +impl AsBuffer for PyByteArray { + fn as_buffer(zelf: &PyObjectView, _vm: &VirtualMachine) -> PyResult { + Ok(PyBuffer::new( + zelf.to_owned().into_object(), + BufferDescriptor::simple(zelf.len(), false), + &BUFFER_METHODS, + )) } } @@ -746,7 +751,7 @@ impl<'a> BufferResizeGuard<'a> for PyByteArray { fn try_resizable(&'a self, vm: &VirtualMachine) -> PyResult { let w = self.inner.upgradable_read(); - if self.exports.load() == 0 { + if self.exports.load(Ordering::SeqCst) == 0 { Ok(parking_lot::lock_api::RwLockUpgradableReadGuard::upgrade(w)) } else { Err(vm diff --git a/vm/src/builtins/bytes.rs b/vm/src/builtins/bytes.rs index 3f4a45762b..0b0ee52289 100644 --- a/vm/src/builtins/bytes.rs +++ b/vm/src/builtins/bytes.rs @@ -6,24 +6,20 @@ use crate::{ bytes_decode, ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions, ByteInnerSplitOptions, ByteInnerTranslateOptions, DecodeArgs, PyBytesInner, }, - common::hash::PyHash, + common::{hash::PyHash, lock::PyMutex}, function::{ ArgBytesLike, ArgIterable, IntoPyObject, IntoPyResult, OptionalArg, OptionalOption, }, - protocol::{BufferInternal, BufferOptions, PyBuffer, PyIterReturn, PyMappingMethods}, + protocol::{BufferDescriptor, BufferMethods, PyBuffer, PyIterReturn, PyMappingMethods}, types::{ AsBuffer, AsMapping, Callable, Comparable, Constructor, Hashable, IterNext, IterNextIterable, Iterable, PyComparisonOp, Unconstructible, }, utils::Either, - IdProtocol, PyClassImpl, PyComparisonValue, PyContext, PyObject, PyObjectRef, PyRef, PyResult, - PyValue, TryFromBorrowedObject, TypeProtocol, VirtualMachine, + IdProtocol, PyClassImpl, PyComparisonValue, PyContext, PyObject, PyObjectRef, PyObjectView, + PyObjectWrap, PyRef, PyResult, PyValue, TryFromBorrowedObject, TypeProtocol, VirtualMachine, }; use bstr::ByteSlice; -use rustpython_common::{ - borrow::{BorrowedValue, BorrowedValueMut}, - lock::PyMutex, -}; use std::mem::size_of; use std::ops::Deref; @@ -541,35 +537,26 @@ impl PyBytes { } } +static BUFFER_METHODS: BufferMethods = BufferMethods { + obj_bytes: |buffer| buffer.obj_as::().as_bytes().into(), + obj_bytes_mut: |_| panic!(), + release: |_| {}, + retain: |_| {}, +}; + impl AsBuffer for PyBytes { - fn as_buffer(zelf: &crate::PyObjectView, _vm: &VirtualMachine) -> PyResult { + fn as_buffer(zelf: &PyObjectView, _vm: &VirtualMachine) -> PyResult { let buf = PyBuffer::new( - zelf.as_object().to_owned(), - zelf.to_owned(), - BufferOptions { - len: zelf.len(), - ..Default::default() - }, + zelf.to_owned().into_object(), + BufferDescriptor::simple(zelf.len(), true), + &BUFFER_METHODS, ); Ok(buf) } } -impl BufferInternal for PyRef { - fn obj_bytes(&self) -> BorrowedValue<[u8]> { - self.as_bytes().into() - } - - fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> { - unreachable!("bytes is not mutable") - } - - fn release(&self) {} - fn retain(&self) {} -} - impl AsMapping for PyBytes { - fn as_mapping(_zelf: &crate::PyObjectView, _vm: &VirtualMachine) -> PyMappingMethods { + fn as_mapping(_zelf: &PyObjectView, _vm: &VirtualMachine) -> PyMappingMethods { PyMappingMethods { length: Some(Self::length), subscript: Some(Self::subscript), diff --git a/vm/src/builtins/memory.rs b/vm/src/builtins/memory.rs index f72a5ee88e..df52bd76af 100644 --- a/vm/src/builtins/memory.rs +++ b/vm/src/builtins/memory.rs @@ -1,26 +1,26 @@ -use super::{PyBytes, PyBytesRef, PyList, PyListRef, PySlice, PyStr, PyStrRef, PyTypeRef}; +use super::{ + PyBytes, PyBytesRef, PyInt, PyListRef, PySlice, PyStr, PyStrRef, PyTuple, PyTupleRef, PyTypeRef, +}; use crate::common::{ borrow::{BorrowedValue, BorrowedValueMut}, hash::PyHash, lock::OnceCell, - rc::PyRc, }; use crate::{ bytesinner::bytes_to_hex, function::{FuncArgs, IntoPyObject, OptionalArg}, - protocol::{BufferInternal, BufferOptions, PyBuffer, PyMappingMethods}, - sliceable::{wrap_index, SaturatedSlice, SequenceIndex}, + protocol::{BufferDescriptor, BufferMethods, PyBuffer, PyMappingMethods, VecBuffer}, + sliceable::wrap_index, stdlib::pystruct::FormatSpec, types::{AsBuffer, AsMapping, Comparable, Constructor, Hashable, PyComparisonOp}, utils::Either, - IdProtocol, PyClassDef, PyClassImpl, PyComparisonValue, PyContext, PyObject, PyObjectRef, - PyRef, PyResult, PyValue, TryFromBorrowedObject, TryFromObject, TypeProtocol, VirtualMachine, + IdProtocol, PyClassImpl, PyComparisonValue, PyContext, PyObject, PyObjectRef, PyObjectView, + PyObjectWrap, PyRef, PyResult, PyValue, TryFromBorrowedObject, TryFromObject, TypeProtocol, + VirtualMachine, }; use crossbeam_utils::atomic::AtomicCell; use itertools::Itertools; -use num_traits::ToPrimitive; -use std::fmt::Debug; -use std::ops::Deref; +use std::{cmp::Ordering, fmt::Debug, mem::ManuallyDrop, ops::Range}; #[derive(FromArgs)] pub struct PyMemoryViewNewArgs { @@ -30,131 +30,103 @@ pub struct PyMemoryViewNewArgs { #[pyclass(module = false, name = "memoryview")] #[derive(Debug)] pub struct PyMemoryView { - buffer: PyBuffer, + // avoid double release when memoryview had released the buffer before drop + buffer: ManuallyDrop, + // the released memoryview does not mean the buffer is destoryed + // because the possible another memeoryview is viewing from it released: AtomicCell, - // start should always less or equal to the stop - // start and stop pointing to the memory index not slice index - // if length is not zero than [start, stop) + // start does NOT mean the bytes before start will not be visited, + // it means the point we starting to get the absolute position via + // the needle start: usize, - stop: usize, - // step can be negative, means read the memory from stop-1 to start - step: isize, format_spec: FormatSpec, + // memoryview's options could be different from buffer's options + desc: BufferDescriptor, hash: OnceCell, + // exports + // memoryview has no exports count by itself + // instead it relay on the buffer it viewing to maintain the count } impl Constructor for PyMemoryView { type Args = PyMemoryViewNewArgs; fn py_new(cls: PyTypeRef, args: Self::Args, vm: &VirtualMachine) -> PyResult { - let buffer = PyBuffer::try_from_borrowed_object(vm, &args.object)?; - let zelf = PyMemoryView::from_buffer(buffer, vm)?; + let zelf = Self::from_object(&args.object, vm)?; zelf.into_pyresult_with_type(vm, cls) } } #[pyimpl(with(Hashable, Comparable, AsBuffer, AsMapping, Constructor))] impl PyMemoryView { - #[cfg(debug_assertions)] - fn validate(self) -> Self { - let options = &self.buffer.options; - let bytes_len = options.len * options.itemsize; - let buffer_len = self.buffer.internal.obj_bytes().len(); - let t1 = self.stop - self.start == bytes_len; - let t2 = buffer_len >= self.stop; - let t3 = buffer_len >= self.start + bytes_len; - assert!(t1); - assert!(t2); - assert!(t3); - self - } - #[cfg(not(debug_assertions))] - fn validate(self) -> Self { - self - } - fn parse_format(format: &str, vm: &VirtualMachine) -> PyResult { FormatSpec::parse(format.as_bytes(), vm) } + /// this should be the main entrence to create the memoryview + /// to avoid the chained memoryview + pub fn from_object(obj: &PyObject, vm: &VirtualMachine) -> PyResult { + if let Some(other) = obj.payload::() { + Ok(other.new_view()) + } else { + let buffer = PyBuffer::try_from_borrowed_object(vm, obj)?; + PyMemoryView::from_buffer(buffer, vm) + } + } + + /// don't use this function to create the memeoryview if the buffer is exporting + /// via another memoryview, use PyMemoryView::new_view() or PyMemoryView::from_object + /// to reduce the chain pub fn from_buffer(buffer: PyBuffer, vm: &VirtualMachine) -> PyResult { // when we get a buffer means the buffered object is size locked // so we can assume the buffer's options will never change as long // as memoryview is still alive - let options = &buffer.options; - let stop = options.len * options.itemsize; - let format_spec = Self::parse_format(&options.format, vm)?; + let format_spec = Self::parse_format(&buffer.desc.format, vm)?; + let desc = buffer.desc.clone(); + Ok(PyMemoryView { - buffer, + buffer: ManuallyDrop::new(buffer), released: AtomicCell::new(false), start: 0, - stop, - step: 1, format_spec, + desc, hash: OnceCell::new(), - } - .validate()) + }) } + /// don't use this function to create the memeoryview if the buffer is exporting + /// via another memoryview, use PyMemoryView::new_view() or PyMemoryView::from_object + /// to reduce the chain pub fn from_buffer_range( - mut buffer: PyBuffer, - range: std::ops::Range, + buffer: PyBuffer, + range: Range, vm: &VirtualMachine, ) -> PyResult { - let itemsize = buffer.options.itemsize; - let format_spec = Self::parse_format(&buffer.options.format, vm)?; - buffer.options.len = range.len(); - Ok(PyMemoryView { - buffer, - released: AtomicCell::new(false), - start: range.start * itemsize, - stop: range.end * itemsize, - step: 1, - format_spec, - hash: OnceCell::new(), - } - .validate()) - } - - fn as_contiguous(&self) -> Option> { - if !self.buffer.options.contiguous { - return None; - } - Some(self.obj_bytes()) - } + let mut zelf = Self::from_buffer(buffer, vm)?; - fn as_contiguous_mut(&self) -> Option> { - if !self.buffer.options.contiguous { - return None; - } - Some(self.obj_bytes_mut()) + zelf.init_range(range, 0); + zelf.init_len(); + Ok(zelf) } - pub fn try_bytes(&self, vm: &VirtualMachine, f: F) -> PyResult - where - F: FnOnce(&[u8]) -> R, - { - self.try_not_released(vm)?; - self.as_contiguous().map(|x| f(&*x)).ok_or_else(|| { - vm.new_type_error("non-contiguous memoryview is not a bytes-like object".to_owned()) - }) + /// this should be the only way to create a memroyview from another memoryview + pub fn new_view(&self) -> Self { + let zelf = PyMemoryView { + buffer: self.buffer.clone(), + released: AtomicCell::new(false), + start: self.start, + format_spec: self.format_spec.clone(), + desc: self.desc.clone(), + hash: OnceCell::new(), + }; + zelf.buffer.retain(); + zelf } #[pymethod] - fn release(&self) { - self._release(); - } - - fn _release(&self) { - // avoid double release + pub fn release(&self) { if self.released.compare_exchange(false, true).is_ok() { - unsafe { - // SAFETY: this branch is only once accessible form _release and guarded by AtomicCell released - let buffer: &std::cell::UnsafeCell = std::mem::transmute(&self.buffer); - let buffer = &mut *buffer.get(); - let internal = std::mem::replace(&mut buffer.internal, PyRc::new(Released)); - internal.release(); - } + self.buffer.release(); } } @@ -173,44 +145,81 @@ impl PyMemoryView { #[pyproperty] fn nbytes(&self, vm: &VirtualMachine) -> PyResult { - self.try_not_released(vm) - .map(|_| self.buffer.options.len * self.buffer.options.itemsize) + self.try_not_released(vm).map(|_| self.desc.len) } #[pyproperty] fn readonly(&self, vm: &VirtualMachine) -> PyResult { - self.try_not_released(vm) - .map(|_| self.buffer.options.readonly) + self.try_not_released(vm).map(|_| self.desc.readonly) } #[pyproperty] fn itemsize(&self, vm: &VirtualMachine) -> PyResult { - self.try_not_released(vm) - .map(|_| self.buffer.options.itemsize) + self.try_not_released(vm).map(|_| self.desc.itemsize) } #[pyproperty] fn ndim(&self, vm: &VirtualMachine) -> PyResult { - self.try_not_released(vm).map(|_| self.buffer.options.ndim) + self.try_not_released(vm).map(|_| self.desc.ndim()) } - // TODO #[pyproperty] - fn shape(&self, vm: &VirtualMachine) -> PyResult { - self.try_not_released(vm) - .map(|_| (self.buffer.options.len,).into_pyobject(vm)) + fn shape(&self, vm: &VirtualMachine) -> PyResult { + self.try_not_released(vm)?; + Ok(vm.ctx.new_tuple( + self.desc + .dim_desc + .iter() + .map(|(shape, _, _)| shape.into_pyobject(vm)) + .collect(), + )) + } + + #[pyproperty] + fn strides(&self, vm: &VirtualMachine) -> PyResult { + self.try_not_released(vm)?; + Ok(vm.ctx.new_tuple( + self.desc + .dim_desc + .iter() + .map(|(_, stride, _)| stride.into_pyobject(vm)) + .collect(), + )) } - // TODO #[pyproperty] - fn strides(&self, vm: &VirtualMachine) -> PyResult { - self.try_not_released(vm).map(|_| (0,).into_pyobject(vm)) + fn suboffsets(&self, vm: &VirtualMachine) -> PyResult { + self.try_not_released(vm)?; + Ok(vm.ctx.new_tuple( + self.desc + .dim_desc + .iter() + .map(|(_, _, suboffset)| suboffset.into_pyobject(vm)) + .collect(), + )) } #[pyproperty] fn format(&self, vm: &VirtualMachine) -> PyResult { self.try_not_released(vm) - .map(|_| PyStr::from(self.buffer.options.format.clone())) + .map(|_| PyStr::from(self.desc.format.clone())) + } + + #[pyproperty] + fn contiguous(&self, vm: &VirtualMachine) -> PyResult { + self.try_not_released(vm).map(|_| self.desc.is_contiguous()) + } + + #[pyproperty] + fn c_contiguous(&self, vm: &VirtualMachine) -> PyResult { + self.try_not_released(vm).map(|_| self.desc.is_contiguous()) + } + + #[pyproperty] + fn f_contiguous(&self, vm: &VirtualMachine) -> PyResult { + // TODO: fortain order + self.try_not_released(vm) + .map(|_| self.desc.ndim() <= 1 && self.desc.is_contiguous()) } #[pymethod(magic)] @@ -220,242 +229,134 @@ impl PyMemoryView { #[pymethod(magic)] fn exit(&self, _args: FuncArgs) { - self._release(); - } - - // translate the slice index to memory index - fn get_pos(&self, i: isize) -> Option { - let len = self.buffer.options.len; - let itemsize = self.buffer.options.itemsize; - let i = wrap_index(i, len)?; - let i = if self.step < 0 { - (len - 1 + (self.step as usize) * i) * itemsize - } else { - self.step as usize * i * itemsize - }; - Some(i) + self.release(); } - fn getitem_by_idx(zelf: PyRef, i: isize, vm: &VirtualMachine) -> PyResult { - let i = zelf - .get_pos(i) + fn getitem_by_idx(&self, i: isize, vm: &VirtualMachine) -> PyResult { + if self.desc.ndim() != 1 { + return Err(vm.new_not_implemented_error( + "multi-dimensional sub-views are not implemented".to_owned(), + )); + } + let (shape, stride, suboffset) = self.desc.dim_desc[0]; + let index = wrap_index(i, shape) .ok_or_else(|| vm.new_index_error("index out of range".to_owned()))?; - let bytes = &zelf.obj_bytes()[i..i + zelf.buffer.options.itemsize]; - zelf.format_spec.unpack(bytes, vm).map(|x| { - if x.len() == 1 { - x.fast_getitem(0) - } else { - x.into() - } - }) + let index = index as isize * stride + suboffset; + let pos = (index + self.start as isize) as usize; + self.unpack_single(pos, vm) } - fn getitem_by_slice(zelf: PyRef, slice: PyRef, vm: &VirtualMachine) -> PyResult { - // slicing a memoryview return a new memoryview - let len = zelf.buffer.options.len; - let (range, step, is_negative_step) = - SaturatedSlice::with_slice(&slice, vm)?.adjust_indices(len); - let abs_step = step.unwrap(); - let step = if is_negative_step { - -(abs_step as isize) - } else { - abs_step as isize - }; - let newstep = step * zelf.step; - let itemsize = zelf.buffer.options.itemsize; - - let format_spec = zelf.format_spec.clone(); - - if newstep == 1 { - let newlen = range.end - range.start; - let start = zelf.start + range.start * itemsize; - let stop = zelf.start + range.end * itemsize; - return Ok(PyMemoryView { - buffer: zelf.buffer.clone_with_options(BufferOptions { - len: newlen, - contiguous: true, - ..zelf.buffer.options.clone() - }), - released: AtomicCell::new(false), - start, - stop, - step: 1, - format_spec, - hash: OnceCell::new(), - } - .validate() - .into_object(vm)); - } - - if range.start >= range.end { - return Ok(PyMemoryView { - buffer: zelf.buffer.clone_with_options(BufferOptions { - len: 0, - contiguous: true, - ..zelf.buffer.options.clone() - }), - released: AtomicCell::new(false), - start: range.end, - stop: range.end, - step: 1, - format_spec, - hash: OnceCell::new(), - } - .validate() - .into_object(vm)); - }; - - // overflow is not possible as dividing a positive integer - let newlen = ((range.end - range.start - 1) / abs_step) - .to_usize() - .unwrap() - + 1; - - // newlen will be 0 if step is overflowed - let newstep = newstep.to_isize().unwrap_or(-1); + fn getitem_by_slice(&self, slice: &PySlice, vm: &VirtualMachine) -> PyResult { + let mut other = self.new_view(); + other.init_slice(slice, 0, vm)?; + other.init_len(); - let (start, stop) = if newstep < 0 { - let stop = zelf.stop - range.start * itemsize * zelf.step.abs() as usize; - let start = stop - (newlen - 1) * itemsize * newstep.abs() as usize - itemsize; - (start, stop) - } else { - let start = zelf.start + range.start * itemsize * zelf.step.abs() as usize; - let stop = start + (newlen - 1) * itemsize * newstep.abs() as usize + itemsize; - (start, stop) - }; + Ok(other.into_ref(vm).into_object()) + } - let options = BufferOptions { - len: newlen, - contiguous: false, - ..zelf.buffer.options.clone() - }; - Ok(PyMemoryView { - buffer: zelf.buffer.clone_with_options(options), - released: AtomicCell::new(false), - start, - stop, - step: newstep, - format_spec, - hash: OnceCell::new(), - } - .validate() - .into_object(vm)) + fn getitem_by_multi_idx(&self, indexes: &[isize], vm: &VirtualMachine) -> PyResult { + let pos = self.pos_from_multi_index(indexes, vm)?; + let bytes = self.buffer.obj_bytes(); + format_unpack(&self.format_spec, &bytes[pos..pos + self.desc.itemsize], vm) } #[pymethod(magic)] fn getitem(zelf: PyRef, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult { zelf.try_not_released(vm)?; - match SequenceIndex::try_from_object_for(vm, needle, Self::NAME)? { - SequenceIndex::Int(index) => Self::getitem_by_idx(zelf, index, vm), - SequenceIndex::Slice(slice) => Self::getitem_by_slice(zelf, slice, vm), + if zelf.desc.ndim() == 0 { + // 0-d memoryview can be referenced using mv[...] or mv[()] only + if needle.is(&vm.ctx.ellipsis) { + return Ok(zelf.into_object()); + } + if let Some(tuple) = needle.payload::() { + if tuple.is_empty() { + return zelf.unpack_single(0, vm); + } + } + return Err(vm.new_type_error("invalid indexing of 0-dim memory".to_owned())); + } + + match SubscriptNeedle::try_from_object(vm, needle)? { + SubscriptNeedle::Index(i) => zelf.getitem_by_idx(i, vm), + SubscriptNeedle::Slice(slice) => zelf.getitem_by_slice(&slice, vm), + SubscriptNeedle::MultiIndex(indices) => zelf.getitem_by_multi_idx(&indices, vm), } } - fn setitem_by_idx( - zelf: PyRef, - i: isize, - value: PyObjectRef, - vm: &VirtualMachine, - ) -> PyResult<()> { - let i = zelf - .get_pos(i) + fn setitem_by_idx(&self, i: isize, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + if self.desc.ndim() != 1 { + return Err(vm.new_not_implemented_error("sub-views are not implemented".to_owned())); + } + let (shape, stride, suboffset) = self.desc.dim_desc[0]; + let index = wrap_index(i, shape) .ok_or_else(|| vm.new_index_error("index out of range".to_owned()))?; - let itemsize = zelf.buffer.options.itemsize; - let data = zelf.format_spec.pack(vec![value], vm)?; - zelf.obj_bytes_mut()[i..i + itemsize].copy_from_slice(&data); - Ok(()) + let index = index as isize * stride + suboffset; + let pos = (index + self.start as isize) as usize; + self.pack_single(pos, value, vm) } fn setitem_by_slice( zelf: PyRef, - slice: PyRef, - items: PyObjectRef, + slice: &PySlice, + src: PyObjectRef, vm: &VirtualMachine, ) -> PyResult<()> { - let items = PyBuffer::try_from_object(vm, items)?; - let options = &items.options; - let len = options.len; - let itemsize = options.itemsize; - - if itemsize != zelf.buffer.options.itemsize { - return Err(vm.new_type_error(format!( - "memoryview: invalid type for format '{}'", - zelf.buffer.options.format - ))); - } - - let diff_err = || { - Err(vm.new_value_error( - "memoryview assignment: lvalue and rvalue have different structures".to_owned(), - )) - }; - - if options.format != zelf.buffer.options.format { - return diff_err(); + if zelf.desc.ndim() != 1 { + return Err(vm.new_not_implemented_error("sub-view are not implemented".to_owned())); } - let (range, step, is_negative_step) = - SaturatedSlice::with_slice(&slice, vm)?.adjust_indices(zelf.buffer.options.len); - - let bytes = items.to_contiguous(); - assert_eq!(bytes.len(), len * itemsize); - - if !is_negative_step && step == Some(1) { - if range.end - range.start != len { - return diff_err(); - } - - if let Some(mut buffer) = zelf.as_contiguous_mut() { - buffer[range.start * itemsize..range.end * itemsize].copy_from_slice(&bytes); - return Ok(()); - } - } + let mut dest = zelf.new_view(); + dest.init_slice(slice, 0, vm)?; + dest.init_len(); - if let Some(step) = step { - let slicelen = if range.end > range.start { - (range.end - range.start - 1) / step + 1 + if zelf.is(&src) { + return if !is_equiv_structure(&zelf.desc, &dest.desc) { + Err(vm.new_value_error( + "memoryview assigment: lvalue and rvalue have different structures".to_owned(), + )) } else { - 0 + // assign self[:] to self + Ok(()) }; + }; - if slicelen != len { - return diff_err(); + let src = if let Some(src) = src.downcast_ref::() { + if zelf.buffer.obj.is(&src.buffer.obj) { + src.to_contiguous(vm) + } else { + AsBuffer::as_buffer(src, vm)? } + } else { + PyBuffer::try_from_object(vm, src)? + }; - let indexes = if is_negative_step { - itertools::Either::Left(range.rev().step_by(step)) - } else { - itertools::Either::Right(range.step_by(step)) - }; + if !is_equiv_structure(&src.desc, &dest.desc) { + return Err(vm.new_value_error( + "memoryview assigment: lvalue and rvalue have different structures".to_owned(), + )); + } - let item_index = (0..len).step_by(itemsize); + let mut bytes_mut = dest.buffer.obj_bytes_mut(); + let src_bytes = src.obj_bytes(); + dest.desc.zip_eq(&src.desc, true, |a_range, b_range| { + let a_range = (a_range.start + dest.start as isize) as usize + ..(a_range.end + dest.start as isize) as usize; + let b_range = b_range.start as usize..b_range.end as usize; + bytes_mut[a_range].copy_from_slice(&src_bytes[b_range]); + false + }); - let mut buffer = zelf.obj_bytes_mut(); + Ok(()) + } - indexes - .map(|i| zelf.get_pos(i as isize).unwrap()) - .zip(item_index) - .for_each(|(i, item_i)| { - buffer[i..i + itemsize].copy_from_slice(&bytes[item_i..item_i + itemsize]); - }); - Ok(()) - } else { - let slicelen = if range.start < range.end { 1 } else { 0 }; - if match len { - 0 => slicelen == 0, - 1 => { - let mut buffer = zelf.obj_bytes_mut(); - let i = zelf.get_pos(range.start as isize).unwrap(); - buffer[i..i + itemsize].copy_from_slice(&bytes); - true - } - _ => false, - } { - Ok(()) - } else { - diff_err() - } - } + fn setitem_by_multi_idx( + &self, + indexes: &[isize], + value: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult<()> { + let pos = self.pos_from_multi_index(indexes, vm)?; + self.pack_single(pos, value, vm) } #[pymethod(magic)] @@ -466,75 +367,203 @@ impl PyMemoryView { vm: &VirtualMachine, ) -> PyResult<()> { zelf.try_not_released(vm)?; - if zelf.buffer.options.readonly { + if zelf.desc.readonly { return Err(vm.new_type_error("cannot modify read-only memory".to_owned())); } - match SequenceIndex::try_from_object_for(vm, needle, Self::NAME)? { - SequenceIndex::Int(index) => Self::setitem_by_idx(zelf, index, value, vm), - SequenceIndex::Slice(slice) => Self::setitem_by_slice(zelf, slice, value, vm), + if value.is(&vm.ctx.none) { + return Err(vm.new_type_error("cannot delete memory".to_owned())); + } + + if zelf.desc.ndim() == 0 { + // TODO: merge branches when we got conditional if let + if needle.is(&vm.ctx.ellipsis) { + return zelf.pack_single(0, value, vm); + } else if let Some(tuple) = needle.payload::() { + if tuple.is_empty() { + return zelf.pack_single(0, value, vm); + } + } + return Err(vm.new_type_error("invalid indexing of 0-dim memory".to_owned())); + } + match SubscriptNeedle::try_from_object(vm, needle)? { + SubscriptNeedle::Index(i) => zelf.setitem_by_idx(i, value, vm), + SubscriptNeedle::Slice(slice) => Self::setitem_by_slice(zelf, &slice, value, vm), + SubscriptNeedle::MultiIndex(indices) => zelf.setitem_by_multi_idx(&indices, value, vm), } } - #[pymethod(magic)] - fn len(&self, vm: &VirtualMachine) -> PyResult { - self.try_not_released(vm).map(|_| self.buffer.options.len) + fn pack_single(&self, pos: usize, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + let mut bytes = self.buffer.obj_bytes_mut(); + // TODO: Optimize + let data = self.format_spec.pack(vec![value], vm).map_err(|_| { + vm.new_type_error(format!( + "memoryview: invalid type for format '{}'", + &self.desc.format + )) + })?; + bytes[pos..pos + self.desc.itemsize].copy_from_slice(&data); + Ok(()) } - fn to_bytes_vec(zelf: &crate::PyObjectView) -> Vec { - if let Some(bytes) = zelf.as_contiguous() { - bytes.to_vec() - } else { - let bytes = &*zelf.obj_bytes(); - let len = zelf.buffer.options.len; - let itemsize = zelf.buffer.options.itemsize; - (0..len) - .map(|i| zelf.get_pos(i as isize).unwrap()) - .flat_map(|i| bytes[i..i + itemsize].to_vec()) - .collect() + fn unpack_single(&self, pos: usize, vm: &VirtualMachine) -> PyResult { + let bytes = self.buffer.obj_bytes(); + // TODO: Optimize + self.format_spec + .unpack(&bytes[pos..pos + self.desc.itemsize], vm) + .map(|x| { + if x.len() == 1 { + x.fast_getitem(0) + } else { + x.into() + } + }) + } + + fn pos_from_multi_index(&self, indexes: &[isize], vm: &VirtualMachine) -> PyResult { + match indexes.len().cmp(&self.desc.ndim()) { + Ordering::Less => { + return Err(vm.new_not_implemented_error("sub-views are not implemented".to_owned())) + } + Ordering::Greater => { + return Err(vm.new_type_error(format!( + "cannot index {}-dimension view with {}-element tuple", + self.desc.ndim(), + indexes.len() + ))) + } + Ordering::Equal => (), } + + let pos = self.desc.position(indexes, vm)?; + let pos = (pos + self.start as isize) as usize; + Ok(pos) } - #[pymethod] - fn tobytes(zelf: PyRef, vm: &VirtualMachine) -> PyResult { - zelf.try_not_released(vm)?; - Ok(PyBytes::from(Self::to_bytes_vec(&zelf)).into_ref(vm)) + fn init_len(&mut self) { + let product: usize = self.desc.dim_desc.iter().map(|x| x.0).product(); + self.desc.len = product * self.desc.itemsize; + } + + fn init_range(&mut self, range: Range, dim: usize) { + let (shape, stride, _) = self.desc.dim_desc[dim]; + debug_assert!(shape >= range.len()); + + let mut is_adjusted = false; + for (_, _, suboffset) in self.desc.dim_desc.iter_mut().rev() { + if *suboffset != 0 { + *suboffset += stride * range.start as isize; + is_adjusted = true; + break; + } + } + if !is_adjusted { + // no suboffset setted, stride must be positive + self.start += stride as usize * range.start; + } + let newlen = range.len(); + self.desc.dim_desc[dim].0 = newlen; + } + + fn init_slice(&mut self, slice: &PySlice, dim: usize, vm: &VirtualMachine) -> PyResult<()> { + let (shape, stride, _) = self.desc.dim_desc[dim]; + let slice = slice.to_saturated(vm)?; + let (range, step, slicelen) = slice.adjust_indices(shape); + + let mut is_adjusted_suboffset = false; + for (_, _, suboffset) in self.desc.dim_desc.iter_mut().rev() { + if *suboffset != 0 { + *suboffset += stride * range.start as isize; + is_adjusted_suboffset = true; + break; + } + } + if !is_adjusted_suboffset { + // no suboffset setted, stride must be positive + self.start += stride as usize + * if step.is_negative() { + range.end - 1 + } else { + range.start + }; + } + self.desc.dim_desc[dim].0 = slicelen; + self.desc.dim_desc[dim].1 *= step; + + Ok(()) + } + + /// return the length of the first dimention + #[pymethod(magic)] + fn len(&self, vm: &VirtualMachine) -> PyResult { + self.try_not_released(vm)?; + Ok(if self.desc.ndim() == 0 { + 1 + } else { + // shape for dim[0] + self.desc.dim_desc[0].0 + }) } #[pymethod] - fn tolist(zelf: PyRef, vm: &VirtualMachine) -> PyResult> { - zelf.try_not_released(vm)?; + fn tobytes(&self, vm: &VirtualMachine) -> PyResult { + self.try_not_released(vm)?; + let mut v = vec![]; + self.append_to(&mut v); + Ok(PyBytes::from(v).into_ref(vm)) + } - let bytes = &*zelf.obj_bytes(); - let elements: Vec = (0..zelf.buffer.options.len) - .map(|i| zelf.get_pos(i as isize).unwrap()) - .map(|i| { - format_unpack( - &zelf.format_spec, - &bytes[i..i + zelf.buffer.options.itemsize], - vm, - ) - }) - .try_collect()?; + fn _to_list( + &self, + bytes: &[u8], + mut index: isize, + dim: usize, + vm: &VirtualMachine, + ) -> PyResult { + let (shape, stride, suboffset) = self.desc.dim_desc[dim]; + if dim + 1 == self.desc.ndim() { + let mut v = Vec::with_capacity(shape); + for _ in 0..shape { + let pos = index + suboffset; + let pos = (pos + self.start as isize) as usize; + let obj = + format_unpack(&self.format_spec, &bytes[pos..pos + self.desc.itemsize], vm)?; + v.push(obj); + index += stride; + } + return Ok(vm.ctx.new_list(v)); + } - Ok(elements) + let mut v = Vec::with_capacity(shape); + for _ in 0..shape { + let obj = self + ._to_list(bytes, index + suboffset, dim + 1, vm)? + .into_object(); + v.push(obj); + index += stride; + } + Ok(vm.ctx.new_list(v)) } #[pymethod] - fn toreadonly(zelf: PyRef, vm: &VirtualMachine) -> PyResult> { - zelf.try_not_released(vm)?; - let buffer = zelf.buffer.clone_with_options(BufferOptions { - readonly: true, - ..zelf.buffer.options.clone() - }); - Ok(PyMemoryView { - buffer, - released: AtomicCell::new(false), - format_spec: zelf.format_spec.clone(), - hash: OnceCell::new(), - ..**zelf + fn tolist(&self, vm: &VirtualMachine) -> PyResult { + self.try_not_released(vm)?; + let bytes = self.buffer.obj_bytes(); + if self.desc.ndim() == 0 { + return Ok(vm.ctx.new_list(vec![format_unpack( + &self.format_spec, + &bytes[..self.desc.itemsize], + vm, + )?])); } - .validate() - .into_ref(vm)) + self._to_list(&bytes, 0, 0, vm) + } + + #[pymethod] + fn toreadonly(&self, vm: &VirtualMachine) -> PyResult> { + self.try_not_released(vm)?; + let mut other = self.new_view(); + other.desc.readonly = true; + Ok(other.into_ref(vm)) } #[pymethod(magic)] @@ -548,63 +577,109 @@ impl PyMemoryView { #[pymethod] fn hex( - zelf: PyRef, + &self, sep: OptionalArg>, bytes_per_sep: OptionalArg, vm: &VirtualMachine, ) -> PyResult { - zelf.try_not_released(vm)?; - let guard; - let vec; - let bytes = match zelf.as_contiguous() { - Some(bytes) => { - guard = bytes; - &*guard - } - None => { - vec = Self::to_bytes_vec(&zelf); - vec.as_slice() - } - }; - - bytes_to_hex(bytes, sep, bytes_per_sep, vm) + self.try_not_released(vm)?; + self.contiguous_or_collect(|x| bytes_to_hex(x, sep, bytes_per_sep, vm)) } - // TODO: support cast shape - #[pymethod] - fn cast(zelf: PyRef, format: PyStrRef, vm: &VirtualMachine) -> PyResult> { - zelf.try_not_released(vm)?; - if !zelf.buffer.options.contiguous { - return Err(vm.new_type_error( - "memoryview: casts are restricted to C-contiguous views".to_owned(), - )); - } - + fn cast_to_1d(&self, format: PyStrRef, vm: &VirtualMachine) -> PyResult { let format_spec = Self::parse_format(format.as_str(), vm)?; let itemsize = format_spec.size(); - let bytelen = zelf.buffer.options.len * zelf.buffer.options.itemsize; - - if bytelen % itemsize != 0 { + if self.desc.len % itemsize != 0 { return Err( vm.new_type_error("memoryview: length is not a multiple of itemsize".to_owned()) ); } - let buffer = zelf.buffer.clone_with_options(BufferOptions { - itemsize, - len: bytelen / itemsize, - format: format.to_string().into(), - ..zelf.buffer.options.clone() - }); - Ok(PyMemoryView { - buffer, + Ok(Self { + buffer: self.buffer.clone(), released: AtomicCell::new(false), + start: self.start, format_spec, + desc: BufferDescriptor { + len: self.desc.len, + readonly: self.desc.readonly, + itemsize, + format: format.to_string().into(), + dim_desc: vec![(self.desc.len / itemsize, itemsize as isize, 0)], + }, hash: OnceCell::new(), - ..**zelf + }) + } + + #[pymethod] + fn cast(&self, args: CastArgs, vm: &VirtualMachine) -> PyResult> { + self.try_not_released(vm)?; + if !self.desc.is_contiguous() { + return Err(vm.new_type_error( + "memoryview: casts are restricted to C-contiguous views".to_owned(), + )); + } + + let CastArgs { format, shape } = args; + + if let OptionalArg::Present(shape) = shape { + if self.desc.is_zero_in_shape() { + return Err(vm.new_type_error( + "memoryview: cannot cast view with zeros in shape or strides".to_owned(), + )); + } + + let shape_vec = shape.borrow_vec(); + let shape_ndim = shape_vec.len(); + // TODO: MAX_NDIM + if self.desc.ndim() != 1 && shape_ndim != 1 { + return Err( + vm.new_type_error("memoryview: cast must be 1D -> ND or ND -> 1D".to_owned()) + ); + } + + let mut other = self.cast_to_1d(format, vm)?; + let itemsize = other.desc.itemsize; + + // 0 ndim is single item + if shape_ndim == 0 { + other.desc.dim_desc = vec![]; + other.desc.len = itemsize; + return Ok(other.into_ref(vm)); + } + + let mut product_shape = itemsize; + let mut dim_descriptor = Vec::with_capacity(shape_ndim); + + for x in shape_vec.iter() { + let x = usize::try_from_borrowed_object(vm, x)?; + + if x > isize::MAX as usize / product_shape { + return Err(vm.new_value_error( + "memoryview.cast(): product(shape) > SSIZE_MAX".to_owned(), + )); + } + product_shape *= x; + dim_descriptor.push((x, 0, 0)); + } + + 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; + } + + if product_shape != other.desc.len { + return Err(vm.new_type_error( + "memoryview: product(shape) * itemsize != buffer size".to_owned(), + )); + } + + other.desc.dim_desc = dim_descriptor; + + Ok(other.into_ref(vm)) + } else { + Ok(self.cast_to_1d(format, vm)?.into_ref(vm)) } - .validate() - .into_ref(vm)) } fn eq( @@ -619,54 +694,62 @@ impl PyMemoryView { return Ok(false); } + if let Some(other) = other.payload::() { + if other.released.load() { + return Ok(false); + } + } + let other = match PyBuffer::try_from_borrowed_object(vm, other) { Ok(buf) => buf, Err(_) => return Ok(false), }; - let a_options = &zelf.buffer.options; - let b_options = &other.options; - - if a_options.len != b_options.len - || a_options.ndim != b_options.ndim - || a_options.shape != b_options.shape - { + if !is_equiv_shape(&zelf.desc, &other.desc) { return Ok(false); } - let a_guard; - let a_vec; - let a = match zelf.as_contiguous() { - Some(bytes) => { - a_guard = bytes; - &*a_guard - } - None => { - a_vec = Self::to_bytes_vec(zelf); - a_vec.as_slice() - } - }; - let b_guard; - let b_vec; - let b = match other.as_contiguous() { - Some(bytes) => { - b_guard = bytes; - &*b_guard - } - None => { - b_vec = other.to_contiguous(); - b_vec.as_slice() - } - }; - - if a_options.format == b_options.format { - Ok(a == b) - } else { - let a_list = unpack_bytes_seq_to_list(a, &a_options.format, vm)?; - let b_list = unpack_bytes_seq_to_list(b, &b_options.format, vm)?; + let a_itemsize = zelf.desc.itemsize; + let b_itemsize = other.desc.itemsize; + let a_format_spec = &zelf.format_spec; + let b_format_spec = &Self::parse_format(&other.desc.format, vm)?; - Ok(vm.bool_eq(a_list.as_object(), b_list.as_object())?) + if zelf.desc.ndim() == 0 { + let a_val = format_unpack(a_format_spec, &zelf.buffer.obj_bytes()[..a_itemsize], vm)?; + let b_val = format_unpack(b_format_spec, &other.obj_bytes()[..b_itemsize], vm)?; + return vm.bool_eq(&a_val, &b_val); } + + // TODO: optimize cmp by format + let mut ret = Ok(true); + let a_bytes = zelf.buffer.obj_bytes(); + let b_bytes = other.obj_bytes(); + zelf.desc.zip_eq(&other.desc, false, |a_range, b_range| { + let a_range = (a_range.start + zelf.start as isize) as usize + ..(a_range.end + zelf.start as isize) as usize; + let b_range = b_range.start as usize..b_range.end as usize; + let a_val = match format_unpack(a_format_spec, &a_bytes[a_range], vm) { + Ok(val) => val, + Err(e) => { + ret = Err(e); + return true; + } + }; + let b_val = match format_unpack(b_format_spec, &b_bytes[b_range], vm) { + Ok(val) => val, + Err(e) => { + ret = Err(e); + return true; + } + }; + ret = vm.bool_eq(&a_val, &b_val); + if let Ok(b) = ret { + !b + } else { + true + } + }); + ret } #[pymethod(magic)] @@ -678,70 +761,185 @@ impl PyMemoryView { fn reduce(_zelf: PyRef, vm: &VirtualMachine) -> PyResult { Err(vm.new_type_error("cannot pickle 'memoryview' object".to_owned())) } -} -impl AsBuffer for PyMemoryView { - fn as_buffer(zelf: &crate::PyObjectView, vm: &VirtualMachine) -> PyResult { - if zelf.released.load() { - Err(vm.new_value_error("operation forbidden on released memoryview object".to_owned())) + fn obj_bytes(&self) -> BorrowedValue<[u8]> { + if self.desc.is_contiguous() { + BorrowedValue::map(self.buffer.obj_bytes(), |x| { + &x[self.start..self.start + self.desc.len] + }) } else { - Ok(PyBuffer::new( - zelf.as_object().to_owned(), - zelf.to_owned(), - zelf.buffer.options.clone(), - )) + BorrowedValue::map(self.buffer.obj_bytes(), |x| &x[self.start..]) } } -} -#[derive(Debug)] -struct Released; -impl BufferInternal for Released { - fn obj_bytes(&self) -> BorrowedValue<[u8]> { - panic!(); + fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> { + if self.desc.is_contiguous() { + BorrowedValueMut::map(self.buffer.obj_bytes_mut(), |x| { + &mut x[self.start..self.start + self.desc.len] + }) + } else { + BorrowedValueMut::map(self.buffer.obj_bytes_mut(), |x| &mut x[self.start..]) + } } - fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> { - panic!(); + fn as_contiguous(&self) -> Option> { + self.desc.is_contiguous().then(|| { + BorrowedValue::map(self.buffer.obj_bytes(), |x| { + &x[self.start..self.start + self.desc.len] + }) + }) } - fn release(&self) {} + fn _as_contiguous_mut(&self) -> Option> { + self.desc.is_contiguous().then(|| { + BorrowedValueMut::map(self.buffer.obj_bytes_mut(), |x| { + &mut x[self.start..self.start + self.desc.len] + }) + }) + } - fn retain(&self) {} -} + fn append_to(&self, buf: &mut Vec) { + if let Some(bytes) = self.as_contiguous() { + buf.extend_from_slice(&bytes); + } else { + buf.reserve(self.desc.len); + let bytes = &*self.buffer.obj_bytes(); + self.desc.for_each_segment(true, |range| { + let start = (range.start + self.start as isize) as usize; + let end = (range.end + self.start as isize) as usize; + buf.extend_from_slice(&bytes[start..end]); + }) + } + } -impl BufferInternal for PyMemoryView { - // NOTE: This impl maybe is anti-pattern. Only used for internal usage. - fn obj_bytes(&self) -> BorrowedValue<[u8]> { - BorrowedValue::map(self.buffer.internal.obj_bytes(), |x| { - &x[self.start..self.stop] - }) + fn contiguous_or_collect R>(&self, f: F) -> R { + let borrowed; + let mut collected; + let v = if let Some(bytes) = self.as_contiguous() { + borrowed = bytes; + &*borrowed + } else { + collected = vec![]; + self.append_to(&mut collected); + &collected + }; + f(v) } - fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> { - BorrowedValueMut::map(self.buffer.internal.obj_bytes_mut(), |x| { - &mut x[self.start..self.stop] - }) + /// clone data from memoryview + /// keep the shape, convert to contiguous + pub fn to_contiguous(&self, vm: &VirtualMachine) -> PyBuffer { + let mut data = vec![]; + self.append_to(&mut data); + + if self.desc.ndim() == 0 { + return VecBuffer::from(data) + .into_ref(vm) + .into_pybuffer_with_descriptor(self.desc.clone()); + } + + let mut dim_desc = self.desc.dim_desc.clone(); + dim_desc.last_mut().unwrap().1 = self.desc.itemsize as isize; + dim_desc.last_mut().unwrap().2 = 0; + for i in (0..dim_desc.len() - 1).rev() { + dim_desc[i].1 = dim_desc[i + 1].1 * dim_desc[i + 1].0 as isize; + dim_desc[i].2 = 0; + } + + let desc = BufferDescriptor { + len: self.desc.len, + readonly: self.desc.readonly, + itemsize: self.desc.itemsize, + format: self.desc.format.clone(), + dim_desc, + }; + + VecBuffer::from(data) + .into_ref(vm) + .into_pybuffer_with_descriptor(desc) } +} - fn release(&self) {} +#[derive(FromArgs)] +struct CastArgs { + #[pyarg(any)] + format: PyStrRef, + #[pyarg(any, optional)] + shape: OptionalArg, +} - fn retain(&self) {} +enum SubscriptNeedle { + Index(isize), + Slice(PyRef), + MultiIndex(Vec), + // MultiSlice(Vec), } -impl BufferInternal for PyRef { - fn obj_bytes(&self) -> BorrowedValue<[u8]> { - self.deref().obj_bytes() +impl TryFromObject for SubscriptNeedle { + fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult { + // TODO: number protocol + if let Some(i) = obj.payload::() { + Ok(Self::Index(i.try_to_primitive(vm)?)) + } else if obj.payload_is::() { + Ok(Self::Slice(unsafe { obj.downcast_unchecked::() })) + } else if let Ok(i) = vm.to_index(&obj) { + Ok(Self::Index(i.try_to_primitive(vm)?)) + } else { + if let Some(tuple) = obj.payload::() { + let tuple = tuple.as_slice(); + if tuple.iter().all(|x| x.payload_is::()) { + let v = tuple + .iter() + .map(|x| { + unsafe { x.downcast_unchecked_ref::() } + .try_to_primitive::(vm) + }) + .try_collect()?; + return Ok(Self::MultiIndex(v)); + } else if tuple.iter().all(|x| x.payload_is::()) { + return Err(vm.new_not_implemented_error( + "multi-dimensional slicing is not implemented".to_owned(), + )); + } + } + Err(vm.new_type_error("memoryview: invalid slice key".to_owned())) + } } - fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> { - self.deref().obj_bytes_mut() +} + +static BUFFER_METHODS: BufferMethods = BufferMethods { + obj_bytes: |buffer| buffer.obj_as::().obj_bytes(), + obj_bytes_mut: |buffer| buffer.obj_as::().obj_bytes_mut(), + release: |buffer| buffer.obj_as::().buffer.release(), + retain: |buffer| buffer.obj_as::().buffer.retain(), +}; + +impl AsBuffer for PyMemoryView { + fn as_buffer(zelf: &PyObjectView, vm: &VirtualMachine) -> PyResult { + if zelf.released.load() { + Err(vm.new_value_error("operation forbidden on released memoryview object".to_owned())) + } else { + Ok(PyBuffer::new( + zelf.to_owned().into_object(), + zelf.desc.clone(), + &BUFFER_METHODS, + )) + } + } +} + +impl Drop for PyMemoryView { + fn drop(&mut self) { + if self.released.load() { + unsafe { self.buffer.drop_without_release() }; + } else { + unsafe { ManuallyDrop::drop(&mut self.buffer) }; + } } - fn release(&self) {} - fn retain(&self) {} } impl AsMapping for PyMemoryView { - fn as_mapping(_zelf: &crate::PyObjectView, _vm: &VirtualMachine) -> PyMappingMethods { + fn as_mapping(_zelf: &PyObjectView, _vm: &VirtualMachine) -> PyMappingMethods { PyMappingMethods { length: Some(Self::length), subscript: Some(Self::subscript), @@ -802,24 +1000,12 @@ impl Hashable for PyMemoryView { zelf.hash .get_or_try_init(|| { zelf.try_not_released(vm)?; - if !zelf.buffer.options.readonly { + if !zelf.desc.readonly { return Err( vm.new_value_error("cannot hash writable memoryview object".to_owned()) ); } - let guard; - let vec; - let bytes = match zelf.as_contiguous() { - Some(bytes) => { - guard = bytes; - &*guard - } - None => { - vec = Self::to_bytes_vec(zelf); - vec.as_slice() - } - }; - Ok(vm.state.hash_secret.hash_bytes(bytes)) + Ok(zelf.contiguous_or_collect(|bytes| vm.state.hash_secret.hash_bytes(bytes))) }) .map(|&x| x) } @@ -849,23 +1035,30 @@ fn format_unpack( }) } -fn unpack_bytes_seq_to_list( - bytes: &[u8], - format: &str, - vm: &VirtualMachine, -) -> PyResult { - let format_spec = PyMemoryView::parse_format(format, vm)?; - let itemsize = format_spec.size(); - - if bytes.len() % itemsize != 0 { - return Err(vm.new_value_error("bytes length not a multiple of item size".to_owned())); +fn is_equiv_shape(a: &BufferDescriptor, b: &BufferDescriptor) -> bool { + if a.ndim() != b.ndim() { + return false; } - let len = bytes.len() / itemsize; + let a_iter = a.dim_desc.iter().map(|x| x.0); + let b_iter = b.dim_desc.iter().map(|x| x.0); + for (a_shape, b_shape) in a_iter.zip(b_iter) { + if a_shape != b_shape { + return false; + } + // if both shape is 0, ignore the rest + if a_shape == 0 { + break; + } + } + true +} - let elements: Vec = (0..len) - .map(|i| format_unpack(&format_spec, &bytes[i..i + itemsize], vm)) - .try_collect()?; +fn is_equiv_format(a: &BufferDescriptor, b: &BufferDescriptor) -> bool { + // TODO: skip @ + a.itemsize == b.itemsize && a.format == b.format +} - Ok(PyList::from(elements).into_ref(vm)) +fn is_equiv_structure(a: &BufferDescriptor, b: &BufferDescriptor) -> bool { + is_equiv_format(a, b) && is_equiv_shape(a, b) } diff --git a/vm/src/builtins/slice.rs b/vm/src/builtins/slice.rs index e1b1350353..99bcec7328 100644 --- a/vm/src/builtins/slice.rs +++ b/vm/src/builtins/slice.rs @@ -9,6 +9,7 @@ use crate::{ use num_bigint::{BigInt, ToBigInt}; use num_traits::{One, Signed, ToPrimitive, Zero}; use std::ops::Range; +use std::option::Option; #[pyclass(module = false, name = "slice")] #[derive(Debug)] @@ -287,53 +288,86 @@ impl SaturatedSlice { // Equivalent to PySlice_AdjustIndices /// Convert for usage in indexing the underlying rust collections. Called *after* /// __index__ has been called on the Slice which might mutate the collection. - pub fn adjust_indices(&self, len: usize) -> (Range, Option, bool) { - // len should always be <= isize::MAX - let ilen = len.to_isize().unwrap_or(isize::MAX); - let (start, stop, step) = (self.start, self.stop, self.step); - let (start, stop, step, is_negative_step) = if step.is_negative() { - ( - if stop == -1 { - ilen.saturating_add(1) - } else { - stop.saturating_add(1) - }, - if start == -1 { - ilen - } else { - start.saturating_add(1) - }, - step.saturating_abs(), - true, - ) + pub fn adjust_indices(&self, len: usize) -> (Range, isize, usize) { + if len == 0 { + return (0..0, self.step, 0); + } + let range = if self.step.is_negative() { + let stop = if self.stop == -1 { + len + } else { + saturate_index(self.stop.saturating_add(1), len) + }; + let start = if self.start == -1 { + len + } else { + saturate_index(self.start.saturating_add(1), len) + }; + stop..start } else { - (start, stop, step, false) + saturate_index(self.start, len)..saturate_index(self.stop, len) }; - let step = step.to_usize(); + let (range, slicelen) = if range.start >= range.end { + (range.start..range.start, 0) + } else { + let slicelen = (range.end - range.start - 1) / self.step.unsigned_abs() + 1; + (range, slicelen) + }; + (range, self.step, slicelen) + } - let range = saturate_index(start, len)..saturate_index(stop, len); - let range = if range.start >= range.end { - range.start..range.start + pub fn iter(&self, len: usize) -> SaturatedSliceIter { + SaturatedSliceIter::new(self, len) + } +} + +pub struct SaturatedSliceIter { + index: isize, + step: isize, + len: usize, +} + +impl SaturatedSliceIter { + pub fn new(slice: &SaturatedSlice, seq_len: usize) -> Self { + let (range, step, len) = slice.adjust_indices(seq_len); + Self::from_adjust_indices(range, step, len) + } + + pub fn from_adjust_indices(range: Range, step: isize, len: usize) -> Self { + let index = if step.is_negative() { + range.end as isize - 1 } else { - // step overflow - if step.is_none() { - if is_negative_step { - (range.end - 1)..range.end - } else { - range.start..(range.start + 1) - } - } else { - range - } + range.start as isize }; - (range, step, is_negative_step) + Self { index, step, len } + } + + pub fn positive_order(mut self) -> Self { + if self.step.is_negative() { + self.index += self.step * self.len.saturating_sub(1) as isize; + self.step = self.step.saturating_abs() + } + self + } +} + +impl Iterator for SaturatedSliceIter { + type Item = usize; + fn next(&mut self) -> Option { + if self.len == 0 { + return None; + } + self.len -= 1; + let ret = self.index as usize; + self.index += self.step; + Some(ret) } } // Go from PyObjectRef to isize w/o overflow error, out of range values are substituted by // isize::MIN or isize::MAX depending on type and value of step. -// Equivalent to PyNumber_AsSsize_t with err equal to None. +// Equivalent to PyEval_SliceIndex. fn to_isize_index(vm: &VirtualMachine, obj: &PyObject) -> PyResult> { if vm.is_none(obj) { return Ok(None); diff --git a/vm/src/cformat.rs b/vm/src/cformat.rs index c166b1fd81..01f67bcc69 100644 --- a/vm/src/cformat.rs +++ b/vm/src/cformat.rs @@ -368,19 +368,7 @@ impl CFormatSpec { } CFormatPreconversor::Str | CFormatPreconversor::Bytes => { if let Ok(buffer) = PyBuffer::try_from_borrowed_object(vm, &obj) { - let guard; - let vec; - let bytes = match buffer.as_contiguous() { - Some(bytes) => { - guard = bytes; - &*guard - } - None => { - vec = buffer.to_contiguous(); - vec.as_slice() - } - }; - Ok(self.format_bytes(bytes)) + Ok(buffer.contiguous_or_collect(|bytes| self.format_bytes(bytes))) } else { let bytes = vm .get_special_method(obj, "__bytes__")? diff --git a/vm/src/function/buffer.rs b/vm/src/function/buffer.rs index 0700d8a312..fd7c20c9dd 100644 --- a/vm/src/function/buffer.rs +++ b/vm/src/function/buffer.rs @@ -40,7 +40,7 @@ impl PyObject { impl ArgBytesLike { pub fn borrow_buf(&self) -> BorrowedValue<'_, [u8]> { - self.0.as_contiguous().unwrap() + unsafe { self.0.contiguous_unchecked() } } pub fn with_ref(&self, f: F) -> R @@ -51,15 +51,11 @@ impl ArgBytesLike { } pub fn len(&self) -> usize { - self.borrow_buf().len() + self.0.desc.len } pub fn is_empty(&self) -> bool { - self.borrow_buf().is_empty() - } - - pub fn to_cow(&self) -> std::borrow::Cow<[u8]> { - self.borrow_buf().to_vec().into() + self.len() == 0 } } @@ -72,7 +68,7 @@ impl From for PyBuffer { impl TryFromBorrowedObject for ArgBytesLike { fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObject) -> PyResult { let buffer = PyBuffer::try_from_borrowed_object(vm, obj)?; - if buffer.options.contiguous { + if buffer.desc.is_contiguous() { Ok(Self(buffer)) } else { Err(vm.new_type_error("non-contiguous buffer is not a bytes-like object".to_owned())) @@ -86,7 +82,7 @@ pub struct ArgMemoryBuffer(PyBuffer); impl ArgMemoryBuffer { pub fn borrow_buf_mut(&self) -> BorrowedValueMut<'_, [u8]> { - self.0.as_contiguous_mut().unwrap() + unsafe { self.0.contiguous_mut_unchecked() } } pub fn with_ref(&self, f: F) -> R @@ -97,11 +93,11 @@ impl ArgMemoryBuffer { } pub fn len(&self) -> usize { - self.borrow_buf_mut().len() + self.0.desc.len } pub fn is_empty(&self) -> bool { - self.borrow_buf_mut().is_empty() + self.len() == 0 } } @@ -114,9 +110,9 @@ impl From for PyBuffer { impl TryFromBorrowedObject for ArgMemoryBuffer { fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObject) -> PyResult { let buffer = PyBuffer::try_from_borrowed_object(vm, obj)?; - if !buffer.options.contiguous { + if !buffer.desc.is_contiguous() { Err(vm.new_type_error("non-contiguous buffer is not a bytes-like object".to_owned())) - } else if buffer.options.readonly { + } else if buffer.desc.readonly { Err(vm.new_type_error("buffer is not a read-write bytes-like object".to_owned())) } else { Ok(Self(buffer)) diff --git a/vm/src/protocol/buffer.rs b/vm/src/protocol/buffer.rs index 1f970ff096..c69b4efc2f 100644 --- a/vm/src/protocol/buffer.rs +++ b/vm/src/protocol/buffer.rs @@ -1,139 +1,435 @@ //! Buffer protocol +//! https://docs.python.org/3/c-api/buffer.html -use crate::common::borrow::{BorrowedValue, BorrowedValueMut}; -use crate::common::rc::PyRc; -use crate::PyThreadingConstraint; -use crate::{PyObject, PyObjectRef, PyResult, TryFromBorrowedObject, TypeProtocol, VirtualMachine}; -use std::{borrow::Cow, fmt::Debug}; - -pub trait BufferInternal: Debug + PyThreadingConstraint { - /// Get the full inner buffer of this memory. You probably want `as_contiguous()`, as - /// `obj_bytes` doesn't take into account the range a memoryview might operate on, among other - /// footguns. - fn obj_bytes(&self) -> BorrowedValue<[u8]>; - /// Get the full inner buffer of this memory, mutably. You probably want - /// `as_contiguous_mut()`, as `obj_bytes` doesn't take into account the range a memoryview - /// might operate on, among other footguns. - fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]>; - fn release(&self); - // not included in PyBuffer protocol itself - fn retain(&self); +use itertools::Itertools; + +use crate::{ + common::{ + borrow::{BorrowedValue, BorrowedValueMut}, + lock::{MapImmutable, PyMutex, PyMutexGuard}, + }, + sliceable::wrap_index, + types::{Constructor, Unconstructible}, + PyObject, PyObjectPayload, PyObjectRef, PyObjectView, PyObjectWrap, PyRef, PyResult, + TryFromBorrowedObject, TypeProtocol, VirtualMachine, +}; +use std::{borrow::Cow, fmt::Debug, ops::Range}; + +#[allow(clippy::type_complexity)] +pub struct BufferMethods { + pub obj_bytes: fn(&PyBuffer) -> BorrowedValue<[u8]>, + pub obj_bytes_mut: fn(&PyBuffer) -> BorrowedValueMut<[u8]>, + pub release: fn(&PyBuffer), + pub retain: fn(&PyBuffer), } -#[derive(Debug)] +impl Debug for BufferMethods { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BufferMethods") + .field("obj_bytes", &(self.obj_bytes as usize)) + .field("obj_bytes_mut", &(self.obj_bytes_mut as usize)) + .field("release", &(self.release as usize)) + .field("retain", &(self.retain as usize)) + .finish() + } +} + +#[derive(Debug, Clone)] pub struct PyBuffer { pub obj: PyObjectRef, - pub options: BufferOptions, - pub(crate) internal: PyRc, + pub desc: BufferDescriptor, + methods: &'static BufferMethods, } impl PyBuffer { - pub fn new( - obj: PyObjectRef, - buffer: impl BufferInternal + 'static, - options: BufferOptions, - ) -> Self { - buffer.retain(); - Self { + pub fn new(obj: PyObjectRef, desc: BufferDescriptor, methods: &'static BufferMethods) -> Self { + let zelf = Self { obj, - options, - internal: PyRc::new(buffer), - } + desc: desc.validate(), + methods, + }; + zelf.retain(); + zelf } + pub fn as_contiguous(&self) -> Option> { - if !self.options.contiguous { - return None; - } - Some(self.internal.obj_bytes()) + self.desc + .is_contiguous() + .then(|| unsafe { self.contiguous_unchecked() }) } pub fn as_contiguous_mut(&self) -> Option> { - if !self.options.contiguous { - return None; - } - Some(self.internal.obj_bytes_mut()) + (!self.desc.readonly && self.desc.is_contiguous()) + .then(|| unsafe { self.contiguous_mut_unchecked() }) } - pub fn to_contiguous(&self) -> Vec { - self.internal.obj_bytes().to_vec() + /// # Safety + /// assume the buffer is contiguous + pub unsafe fn contiguous_unchecked(&self) -> BorrowedValue<[u8]> { + self.obj_bytes() } - pub fn clone_with_options(&self, options: BufferOptions) -> Self { - self.internal.retain(); - Self { - obj: self.obj.clone(), - options, - internal: self.internal.clone(), + /// # Safety + /// assume the buffer is contiguous and writable + pub unsafe fn contiguous_mut_unchecked(&self) -> BorrowedValueMut<[u8]> { + self.obj_bytes_mut() + } + + pub fn append_to(&self, buf: &mut Vec) { + if let Some(bytes) = self.as_contiguous() { + buf.extend_from_slice(&bytes); + } else { + let bytes = &*self.obj_bytes(); + self.desc.for_each_segment(true, |range| { + buf.extend_from_slice(&bytes[range.start as usize..range.end as usize]) + }); } } -} -#[derive(Debug, Clone)] -pub struct BufferOptions { - // buf - pub len: usize, - pub readonly: bool, - pub itemsize: usize, - pub format: Cow<'static, str>, - pub ndim: usize, // TODO: support multiple dimension array - pub shape: Vec, - pub strides: Vec, - // suboffsets + pub fn contiguous_or_collect R>(&self, f: F) -> R { + let borrowed; + let mut collected; + let v = if let Some(bytes) = self.as_contiguous() { + borrowed = bytes; + &*borrowed + } else { + collected = vec![]; + self.append_to(&mut collected); + &collected + }; + f(v) + } - // RustPython fields - pub contiguous: bool, -} + pub fn obj_as(&self) -> &PyObjectView { + unsafe { self.obj.downcast_unchecked_ref() } + } -impl BufferOptions { - pub const DEFAULT: Self = BufferOptions { - len: 0, - readonly: true, - itemsize: 1, - format: Cow::Borrowed("B"), - ndim: 1, - shape: Vec::new(), - strides: Vec::new(), - contiguous: true, - }; -} + pub fn obj_bytes(&self) -> BorrowedValue<[u8]> { + (self.methods.obj_bytes)(self) + } + + pub fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> { + (self.methods.obj_bytes_mut)(self) + } + + pub fn release(&self) { + (self.methods.release)(self) + } -impl Default for BufferOptions { - fn default() -> Self { - Self::DEFAULT + pub fn retain(&self) { + (self.methods.retain)(self) + } + + // drop PyBuffer without calling release + // after this function, the owner should use forget() + // or wrap PyBuffer in the ManaullyDrop to prevent drop() + pub(crate) unsafe fn drop_without_release(&mut self) { + std::ptr::drop_in_place(&mut self.obj); + std::ptr::drop_in_place(&mut self.desc); } } impl TryFromBorrowedObject for PyBuffer { fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObject) -> PyResult { - let obj_cls = obj.class(); - for cls in obj_cls.iter_mro() { - if let Some(f) = cls.slots.as_buffer.as_ref() { - return f(obj, vm); - } + let cls = obj.class(); + if let Some(f) = cls.mro_find_map(|cls| cls.slots.as_buffer) { + return f(obj, vm); } Err(vm.new_type_error(format!( "a bytes-like object is required, not '{}'", - obj_cls.name() + cls.name() ))) } } -// What we actually want to implement is: -// impl Drop for T where T: BufferInternal -// but it is not supported by Rust impl Drop for PyBuffer { fn drop(&mut self) { - self.internal.release(); + self.release(); } } -impl Clone for PyBuffer { - fn clone(&self) -> Self { - self.clone_with_options(self.options.clone()) +#[derive(Debug, Clone)] +pub struct BufferDescriptor { + /// product(shape) * itemsize + /// bytes length, but not the length for obj_bytes() even is contiguous + pub len: usize, + pub readonly: bool, + pub itemsize: usize, + pub format: Cow<'static, str>, + /// (shape, stride, suboffset) for each dimension + pub dim_desc: Vec<(usize, isize, isize)>, + // TODO: flags +} + +impl BufferDescriptor { + pub fn simple(bytes_len: usize, readonly: bool) -> Self { + Self { + len: bytes_len, + readonly, + itemsize: 1, + format: Cow::Borrowed("B"), + dim_desc: vec![(bytes_len, 1, 0)], + } + } + + pub fn format( + bytes_len: usize, + readonly: bool, + itemsize: usize, + format: Cow<'static, str>, + ) -> Self { + Self { + len: bytes_len, + readonly, + itemsize, + format, + dim_desc: vec![(bytes_len / itemsize, itemsize as isize, 0)], + } + } + + #[cfg(debug_assertions)] + pub fn validate(self) -> Self { + assert!(self.itemsize != 0); + assert!(self.ndim() != 0); + let mut shape_product = 1; + for (shape, stride, suboffset) in self.dim_desc.iter().cloned() { + shape_product *= shape; + assert!(suboffset >= 0); + assert!(stride != 0); + } + assert!(shape_product * self.itemsize == self.len); + self + } + + #[cfg(not(debug_assertions))] + pub fn validate(self) -> Self { + self + } + + pub fn ndim(&self) -> usize { + self.dim_desc.len() } + + pub fn is_contiguous(&self) -> bool { + if self.len == 0 { + return true; + } + let mut sd = self.itemsize; + for (shape, stride, _) in self.dim_desc.iter().cloned().rev() { + if shape > 1 && stride != sd as isize { + return false; + } + sd *= shape; + } + true + } + + /// this function do not check the bound + /// panic if indices.len() != ndim + pub fn fast_position(&self, indices: &[usize]) -> isize { + let mut pos = 0; + for (i, (_, stride, suboffset)) in indices + .iter() + .cloned() + .zip_eq(self.dim_desc.iter().cloned()) + { + pos += i as isize * stride + suboffset; + } + pos + } + + /// panic if indices.len() != ndim + pub fn position(&self, indices: &[isize], vm: &VirtualMachine) -> PyResult { + let mut pos = 0; + for (i, (shape, stride, suboffset)) in indices + .iter() + .cloned() + .zip_eq(self.dim_desc.iter().cloned()) + { + let i = wrap_index(i, shape).ok_or_else(|| { + vm.new_index_error(format!("index out of bounds on dimension {}", i)) + })?; + pos += i as isize * stride + suboffset; + } + Ok(pos) + } + + pub fn for_each_segment(&self, try_conti: bool, mut f: F) + where + F: FnMut(Range), + { + if self.ndim() == 0 { + f(0..self.itemsize as isize); + return; + } + if try_conti && self.is_last_dim_contiguous() { + self._for_each_segment::<_, true>(0, 0, &mut f); + } else { + self._for_each_segment::<_, false>(0, 0, &mut f); + } + } + + fn _for_each_segment(&self, mut index: isize, dim: usize, f: &mut F) + where + F: FnMut(Range), + { + let (shape, stride, suboffset) = self.dim_desc[dim]; + if dim + 1 == self.ndim() { + if CONTI { + f(index..index + (shape * self.itemsize) as isize); + } else { + for _ in 0..shape { + let pos = index + suboffset; + f(pos..pos + self.itemsize as isize); + index += stride; + } + } + return; + } + for _ in 0..shape { + self._for_each_segment::(index + suboffset, dim + 1, f); + index += stride; + } + } + + /// zip two BufferDescriptor with the same shape + pub fn zip_eq(&self, other: &Self, try_conti: bool, mut f: F) + where + F: FnMut(Range, Range) -> bool, + { + if self.ndim() == 0 { + f(0..self.itemsize as isize, 0..other.itemsize as isize); + return; + } + if try_conti && self.is_last_dim_contiguous() { + self._zip_eq::<_, true>(other, 0, 0, 0, &mut f); + } else { + self._zip_eq::<_, false>(other, 0, 0, 0, &mut f); + } + } + + fn _zip_eq( + &self, + other: &Self, + mut a_index: isize, + mut b_index: isize, + dim: usize, + f: &mut F, + ) 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); + if dim + 1 == self.ndim() { + if CONTI { + if f( + a_index..a_index + (shape * self.itemsize) as isize, + b_index..b_index + (shape * other.itemsize) as isize, + ) { + return; + } + } else { + for _ in 0..shape { + let a_pos = a_index + a_suboffset; + let b_pos = b_index + b_suboffset; + if f( + a_pos..a_pos + self.itemsize as isize, + b_pos..b_pos + other.itemsize as isize, + ) { + return; + } + a_index += a_stride; + b_index += b_stride; + } + } + return; + } + + for _ in 0..shape { + self._zip_eq::( + other, + a_index + a_suboffset, + b_index + b_suboffset, + dim + 1, + f, + ); + a_index += a_stride; + b_index += b_stride; + } + } + + fn is_last_dim_contiguous(&self) -> bool { + let (_, stride, suboffset) = self.dim_desc[self.ndim() - 1]; + suboffset == 0 && stride == self.itemsize as isize + } + + pub fn is_zero_in_shape(&self) -> bool { + for (shape, _, _) in self.dim_desc.iter().cloned() { + if shape == 0 { + return true; + } + } + false + } + + // TODO: support fortain order } pub trait BufferResizeGuard<'a> { type Resizable: 'a; fn try_resizable(&'a self, vm: &VirtualMachine) -> PyResult; } + +#[pyclass(module = false, name = "vec_buffer")] +#[derive(Debug, PyValue)] +pub struct VecBuffer { + data: PyMutex>, +} + +#[pyimpl(flags(BASETYPE), with(Constructor))] +impl VecBuffer { + pub fn take(&self) -> Vec { + std::mem::take(&mut self.data.lock()) + } +} + +impl From> for VecBuffer { + fn from(data: Vec) -> Self { + Self { + data: PyMutex::new(data), + } + } +} + +impl Unconstructible for VecBuffer {} + +impl PyRef { + pub fn into_pybuffer(self, readonly: bool) -> PyBuffer { + let len = self.data.lock().len(); + PyBuffer::new( + self.into_object(), + BufferDescriptor::simple(len, readonly), + &VEC_BUFFER_METHODS, + ) + } + + pub fn into_pybuffer_with_descriptor(self, desc: BufferDescriptor) -> PyBuffer { + PyBuffer::new(self.into_object(), desc, &VEC_BUFFER_METHODS) + } +} + +static VEC_BUFFER_METHODS: BufferMethods = BufferMethods { + obj_bytes: |buffer| { + PyMutexGuard::map_immutable(buffer.obj_as::().data.lock(), |x| x.as_slice()) + .into() + }, + obj_bytes_mut: |buffer| { + PyMutexGuard::map(buffer.obj_as::().data.lock(), |x| { + x.as_mut_slice() + }) + .into() + }, + release: |_| {}, + retain: |_| {}, +}; diff --git a/vm/src/protocol/mod.rs b/vm/src/protocol/mod.rs index 38c6a547e5..24c88c31da 100644 --- a/vm/src/protocol/mod.rs +++ b/vm/src/protocol/mod.rs @@ -3,6 +3,6 @@ mod iter; mod mapping; mod object; -pub use buffer::{BufferInternal, BufferOptions, BufferResizeGuard, PyBuffer}; +pub use buffer::{BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, VecBuffer}; pub use iter::{PyIter, PyIterIter, PyIterReturn}; pub use mapping::{PyMapping, PyMappingMethods}; diff --git a/vm/src/sliceable.rs b/vm/src/sliceable.rs index 6eab0e5c64..5f7e388cec 100644 --- a/vm/src/sliceable.rs +++ b/vm/src/sliceable.rs @@ -3,8 +3,8 @@ use std::ops::Range; use crate::builtins::int::PyInt; // export through slicable module, not slice. -use crate::builtins::slice::PySlice; pub use crate::builtins::slice::{saturate_index, SaturatedSlice}; +use crate::builtins::slice::{PySlice, SaturatedSliceIter}; use crate::utils::Either; use crate::VirtualMachine; use crate::{PyObjectRef, PyRef, PyResult, TypeProtocol}; @@ -28,55 +28,19 @@ pub trait PySliceableSequenceMut { slice: SaturatedSlice, items: &[Self::Item], ) -> PyResult<()> { - let (range, step, is_negative_step) = slice.adjust_indices(self.as_slice().len()); - if !is_negative_step && step == Some(1) { - return if range.end - range.start == items.len() { - self.do_set_range(range, items); - Ok(()) - } else { - Err(vm.new_buffer_error( - "Existing exports of data: object cannot be re-sized".to_owned(), - )) - }; - } - if let Some(step) = step { - let slicelen = if range.end > range.start { - (range.end - range.start - 1) / step + 1 - } else { - 0 - }; - - if slicelen == items.len() { - let indexes = if is_negative_step { - itertools::Either::Left(range.rev().step_by(step)) - } else { - itertools::Either::Right(range.step_by(step)) - }; - self.do_replace_indexes(indexes, items); - Ok(()) - } else { - Err(vm.new_buffer_error( - "Existing exports of data: object cannot be re-sized".to_owned(), - )) - } + let (range, step, slicelen) = slice.adjust_indices(self.as_slice().len()); + if slicelen != items.len() { + Err(vm + .new_buffer_error("Existing exports of data: object cannot be re-sized".to_owned())) + } else if step == 1 { + self.do_set_range(range, items); + Ok(()) } else { - // edge case, step is too big for usize - // same behaviour as CPython - let slicelen = if range.start < range.end { 1 } else { 0 }; - if match items.len() { - 0 => slicelen == 0, - 1 => { - self.do_set_range(range, items); - true - } - _ => false, - } { - Ok(()) - } else { - Err(vm.new_buffer_error( - "Existing exports of data: object cannot be re-sized".to_owned(), - )) - } + self.do_replace_indexes( + SaturatedSliceIter::from_adjust_indices(range, step, slicelen), + items, + ); + Ok(()) } } @@ -86,82 +50,39 @@ pub trait PySliceableSequenceMut { slice: SaturatedSlice, items: &[Self::Item], ) -> PyResult<()> { - let (range, step, is_negative_step) = slice.adjust_indices(self.as_slice().len()); - if !is_negative_step && step == Some(1) { + let (range, step, slicelen) = slice.adjust_indices(self.as_slice().len()); + if step == 1 { self.do_set_range(range, items); - return Ok(()); - } - if let Some(step) = step { - let slicelen = if range.end > range.start { - (range.end - range.start - 1) / step + 1 - } else { - 0 - }; - - if slicelen == items.len() { - let indexes = if is_negative_step { - itertools::Either::Left(range.rev().step_by(step)) - } else { - itertools::Either::Right(range.step_by(step)) - }; - self.do_replace_indexes(indexes, items); - Ok(()) - } else { - Err(vm.new_value_error(format!( - "attempt to assign sequence of size {} to extended slice of size {}", - items.len(), - slicelen - ))) - } + Ok(()) + } else if slicelen == items.len() { + self.do_replace_indexes( + SaturatedSliceIter::from_adjust_indices(range, step, slicelen), + items, + ); + Ok(()) } else { - // edge case, step is too big for usize - // same behaviour as CPython - let slicelen = if range.start < range.end { 1 } else { 0 }; - if match items.len() { - 0 => slicelen == 0, - 1 => { - self.do_set_range(range, items); - true - } - _ => false, - } { - Ok(()) - } else { - Err(vm.new_value_error(format!( - "attempt to assign sequence of size {} to extended slice of size {}", - items.len(), - slicelen - ))) - } + Err(vm.new_value_error(format!( + "attempt to assign sequence of size {} to extended slice of size {}", + items.len(), + slicelen + ))) } } fn delete_slice(&mut self, _vm: &VirtualMachine, slice: SaturatedSlice) -> PyResult<()> { - let (range, step, is_negative_step) = slice.adjust_indices(self.as_slice().len()); - if range.start >= range.end { - return Ok(()); - } - - if !is_negative_step && step == Some(1) { + let (range, step, slicelen) = slice.adjust_indices(self.as_slice().len()); + if slicelen == 0 { + Ok(()) + } else if step == 1 { self.do_delete_range(range); - return Ok(()); - } - - // step is not negative here - if let Some(step) = step { - let indexes = if is_negative_step { - itertools::Either::Left(range.clone().rev().step_by(step).rev()) - } else { - itertools::Either::Right(range.clone().step_by(step)) - }; - - self.do_delete_indexes(range, indexes); + Ok(()) } else { - // edge case, step is too big for usize - // same behaviour as CPython - self.do_delete_range(range); + self.do_delete_indexes( + range.clone(), + SaturatedSliceIter::from_adjust_indices(range, step, slicelen).positive_order(), + ); + Ok(()) } - Ok(()) } } @@ -237,28 +158,21 @@ pub trait PySliceableSequence { _vm: &VirtualMachine, slice: SaturatedSlice, ) -> PyResult { - let (range, step, is_negative_step) = slice.adjust_indices(self.len()); - if range.start >= range.end { - return Ok(Self::empty()); - } - - if step == Some(1) { - return Ok(if is_negative_step { + let (range, step, slicelen) = slice.adjust_indices(self.len()); + let sliced = if slicelen == 0 { + Self::empty() + } else if step == 1 { + if step.is_negative() { self.do_slice_reverse(range) } else { self.do_slice(range) - }); - } - - if let Some(step) = step { - Ok(if is_negative_step { - self.do_stepped_slice_reverse(range, step) - } else { - self.do_stepped_slice(range, step) - }) + } + } else if step.is_negative() { + self.do_stepped_slice_reverse(range, step.unsigned_abs()) } else { - Ok(self.do_slice(range)) - } + self.do_stepped_slice(range, step.unsigned_abs()) + }; + Ok(sliced) } fn get_item( diff --git a/vm/src/stdlib/builtins.rs b/vm/src/stdlib/builtins.rs index 1c428b7099..4dace44c43 100644 --- a/vm/src/stdlib/builtins.rs +++ b/vm/src/stdlib/builtins.rs @@ -1,7 +1,7 @@ //! Builtin function definitions. //! //! Implements the list of [builtin Python functions](https://docs.python.org/3/library/builtins.html). -use crate::{PyObjectRef, VirtualMachine}; +use crate::{PyClassImpl, PyObjectRef, VirtualMachine}; /// Built-in functions, exceptions, and other objects. /// @@ -917,7 +917,9 @@ pub use builtins::{ascii, print}; pub fn make_module(vm: &VirtualMachine, module: PyObjectRef) { let ctx = &vm.ctx; - builtins::extend_module(vm, &module); + crate::protocol::VecBuffer::make_class(&vm.ctx); + + let _ = builtins::extend_module(vm, &module); let debug_mode: bool = vm.state.settings.optimize == 0; extend_module!(vm, module, { diff --git a/vm/src/stdlib/io.rs b/vm/src/stdlib/io.rs index a191177078..a2174f114e 100644 --- a/vm/src/stdlib/io.rs +++ b/vm/src/stdlib/io.rs @@ -68,29 +68,27 @@ impl TryFromObject for Fildes { mod _io { use super::*; - use crate::common::{ - borrow::{BorrowedValue, BorrowedValueMut}, - lock::{ - PyMappedThreadMutexGuard, PyMutex, PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard, - PyThreadMutex, PyThreadMutexGuard, - }, - rc::PyRc, - }; use crate::{ builtins::{ PyBaseExceptionRef, PyByteArray, PyBytes, PyBytesRef, PyIntRef, PyMemoryView, PyStr, PyStrRef, PyType, PyTypeRef, }, + common::lock::{ + PyMappedThreadMutexGuard, PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard, + PyThreadMutex, PyThreadMutexGuard, + }, function::{ ArgBytesLike, ArgIterable, ArgMemoryBuffer, FuncArgs, IntoPyObject, OptionalArg, OptionalOption, }, - protocol::{BufferInternal, BufferOptions, BufferResizeGuard, PyBuffer, PyIterReturn}, + protocol::{ + BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn, VecBuffer, + }, types::{Constructor, Destructor, IterNext, Iterable}, utils::Either, vm::{ReprGuard, VirtualMachine}, - IdProtocol, PyContext, PyObject, PyObjectRef, PyRef, PyResult, PyValue, StaticType, - TryFromBorrowedObject, TryFromObject, TypeProtocol, + IdProtocol, PyContext, PyObject, PyObjectRef, PyObjectWrap, PyRef, PyResult, PyValue, + StaticType, TryFromBorrowedObject, TryFromObject, TypeProtocol, }; use bstr::ByteSlice; use crossbeam_utils::atomic::AtomicCell; @@ -876,32 +874,20 @@ mod _io { // TODO: loop if write() raises an interrupt vm.call_method(self.raw.as_ref().unwrap(), "write", (memobj,))? } else { - let options = BufferOptions { - len, - ..Default::default() - }; - // TODO: see if we can encapsulate this pattern in a function in memory.rs like - // fn slice_as_memory(s: &[u8], f: impl FnOnce(PyMemoryViewRef) -> R) -> R - let writebuf = PyRc::new(BufferedRawBuffer { - data: std::mem::take(&mut self.buffer).into(), - range: buf_range, - }); - let raw = self.raw.as_ref().unwrap(); - let memobj = PyMemoryView::from_buffer( - PyBuffer { - obj: raw.clone(), - options, - internal: writebuf.clone(), - }, + let v = std::mem::take(&mut self.buffer); + let writebuf = VecBuffer::from(v).into_ref(vm); + let memobj = PyMemoryView::from_buffer_range( + writebuf.clone().into_pybuffer(true), + buf_range, vm, )? .into_ref(vm); // TODO: loop if write() raises an interrupt - let res = vm.call_method(raw, "write", (memobj.clone(),)); + let res = vm.call_method(self.raw.as_ref().unwrap(), "write", (memobj.clone(),)); memobj.release(); - self.buffer = std::mem::take(&mut writebuf.data.lock()); + self.buffer = writebuf.take(); res? }; @@ -1114,23 +1100,10 @@ mod _io { let res = match v { Either::A(v) => { let v = v.unwrap_or(&mut self.buffer); - let options = BufferOptions { - len, - readonly: false, - ..Default::default() - }; - // TODO: see if we can encapsulate this pattern in a function in memory.rs like - // fn slice_as_memory(s: &[u8], f: impl FnOnce(PyMemoryViewRef) -> R) -> R - let readbuf = PyRc::new(BufferedRawBuffer { - data: std::mem::take(v).into(), - range: buf_range, - }); - let memobj = PyMemoryView::from_buffer( - PyBuffer { - obj: vm.ctx.none(), - options, - internal: readbuf.clone(), - }, + let readbuf = VecBuffer::from(std::mem::take(v)).into_ref(vm); + let memobj = PyMemoryView::from_buffer_range( + readbuf.clone().into_pybuffer(false), + buf_range, vm, )? .into_ref(vm); @@ -1140,7 +1113,7 @@ mod _io { vm.call_method(self.raw.as_ref().unwrap(), "readinto", (memobj.clone(),)); memobj.release(); - std::mem::swap(v, &mut readbuf.data.lock()); + std::mem::swap(v, &mut readbuf.take()); res? } @@ -1319,29 +1292,6 @@ mod _io { } } - // this is a bit fancier than what CPython does, but in CPython if you store - // the memoryobj for the buffer until after the BufferedIO is destroyed, you - // can get a use-after-free, so this is a bit safe - #[derive(Debug)] - struct BufferedRawBuffer { - data: PyMutex>, - range: Range, - } - impl BufferInternal for BufferedRawBuffer { - fn obj_bytes(&self) -> BorrowedValue<[u8]> { - BorrowedValue::map(self.data.lock().into(), |data| &data[self.range.clone()]) - } - - fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> { - BorrowedValueMut::map(self.data.lock().into(), |data| { - &mut data[self.range.clone()] - }) - } - - fn release(&self) {} - fn retain(&self) {} - } - pub fn get_offset(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult { let int = vm.to_index(&obj)?; int.as_bigint().try_into().map_err(|_| { @@ -3334,35 +3284,36 @@ mod _io { #[pymethod] fn getbuffer(self, vm: &VirtualMachine) -> PyResult { - let options = BufferOptions { - readonly: false, - len: self.buffer.read().cursor.get_ref().len(), - ..Default::default() - }; - let buffer = PyBuffer::new(self.as_object().to_owned(), self, options); + let len = self.buffer.read().cursor.get_ref().len(); + let buffer = PyBuffer::new( + self.into_object(), + BufferDescriptor::simple(len, false), + &BYTES_IO_BUFFER_METHODS, + ); let view = PyMemoryView::from_buffer(buffer, vm)?; Ok(view) } } - impl BufferInternal for PyRef { - fn obj_bytes(&self) -> BorrowedValue<[u8]> { - PyRwLockReadGuard::map(self.buffer.read(), |x| x.cursor.get_ref().as_slice()).into() - } - - fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> { - PyRwLockWriteGuard::map(self.buffer.write(), |x| x.cursor.get_mut().as_mut_slice()) + static BYTES_IO_BUFFER_METHODS: BufferMethods = BufferMethods { + obj_bytes: |buffer| { + let zelf = buffer.obj_as::(); + PyRwLockReadGuard::map(zelf.buffer.read(), |x| x.cursor.get_ref().as_slice()).into() + }, + obj_bytes_mut: |buffer| { + let zelf = buffer.obj_as::(); + PyRwLockWriteGuard::map(zelf.buffer.write(), |x| x.cursor.get_mut().as_mut_slice()) .into() - } + }, - fn release(&self) { - self.exports.fetch_sub(1); - } + release: |buffer| { + buffer.obj_as::().exports.fetch_sub(1); + }, - fn retain(&self) { - self.exports.fetch_add(1); - } - } + retain: |buffer| { + buffer.obj_as::().exports.fetch_add(1); + }, + }; impl<'a> BufferResizeGuard<'a> for BytesIO { type Resizable = PyRwLockWriteGuard<'a, BufferedIO>; diff --git a/vm/src/stdlib/os.rs b/vm/src/stdlib/os.rs index 63bddbcf01..e83647f5f9 100644 --- a/vm/src/stdlib/os.rs +++ b/vm/src/stdlib/os.rs @@ -191,7 +191,11 @@ impl TryFromObject for PyPathLike { fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult { // path_converter in CPython let obj = match PyBuffer::try_from_borrowed_object(vm, &obj) { - Ok(buffer) => PyBytes::from(buffer.internal.obj_bytes().to_vec()).into_pyobject(vm), + Ok(buffer) => { + let mut bytes = vec![]; + buffer.append_to(&mut bytes); + PyBytes::from(bytes).into_pyobject(vm) + } Err(_) => obj, }; let fs_path = FsPath::try_from(obj, true, vm)?; diff --git a/vm/src/stdlib/sre.rs b/vm/src/stdlib/sre.rs index 7c216b2bf5..c9c8ba286e 100644 --- a/vm/src/stdlib/sre.rs +++ b/vm/src/stdlib/sre.rs @@ -140,31 +140,15 @@ mod _sre { vm: &VirtualMachine, f: F, ) -> PyResult { - let buffer; - let guard; - let vec; - let s; - let str_drive = if self.isbytes { - buffer = PyBuffer::try_from_borrowed_object(vm, &string)?; - let bytes = match buffer.as_contiguous() { - Some(bytes) => { - guard = bytes; - &*guard - } - None => { - vec = buffer.to_contiguous(); - vec.as_slice() - } - }; - StrDrive::Bytes(bytes) + if self.isbytes { + let buffer = PyBuffer::try_from_borrowed_object(vm, &string)?; + buffer.contiguous_or_collect(|bytes| f(StrDrive::Bytes(bytes))) } else { - s = string + let s = string .payload::() .ok_or_else(|| vm.new_type_error("expected string".to_owned()))?; - StrDrive::Str(s.as_str()) - }; - - f(str_drive) + f(StrDrive::Str(s.as_str())) + } } fn with_state PyResult>(