diff --git a/vm/src/builtins/memory.rs b/vm/src/builtins/memory.rs index 4b7bace0e..bd2d3319d 100644 --- a/vm/src/builtins/memory.rs +++ b/vm/src/builtins/memory.rs @@ -9,6 +9,7 @@ use crate::bytesinner::bytes_to_hex; use crate::common::borrow::{BorrowedValue, BorrowedValueMut}; use crate::common::hash::PyHash; use crate::common::lock::OnceCell; +use crate::common::rc::PyRc; use crate::function::OptionalArg; use crate::pyobject::{ Either, IdProtocol, IntoPyObject, PyClassImpl, PyComparisonValue, PyContext, PyObjectRef, @@ -41,12 +42,55 @@ impl BufferRef { pub fn new(buffer: impl Buffer + 'static) -> Self { Self(Box::new(buffer)) } + pub fn into_rcbuf(self) -> RcBuffer { + // move self.0 out of self; BufferRef impls Drop so it's tricky + let this = std::mem::ManuallyDrop::new(self); + let buf_box = unsafe { std::ptr::read(&this.0) }; + RcBuffer(buf_box.into()) + } } impl From> for BufferRef { fn from(buffer: Box) -> Self { BufferRef(buffer) } } +#[derive(Debug, Clone)] +pub struct RcBuffer(PyRc); +impl Deref for RcBuffer { + type Target = dyn Buffer; + fn deref(&self) -> &Self::Target { + self.0.deref() + } +} +impl Drop for RcBuffer { + fn drop(&mut self) { + // check if this is the last rc before the inner buffer gets dropped + if let Some(buf) = PyRc::get_mut(&mut self.0) { + buf.release() + } + } +} +impl Buffer for RcBuffer { + fn get_options(&self) -> BorrowedValue { + self.0.get_options() + } + fn obj_bytes(&self) -> BorrowedValue<[u8]> { + self.0.obj_bytes() + } + fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> { + self.0.obj_bytes_mut() + } + fn release(&self) {} + fn as_contiguous(&self) -> Option> { + self.0.as_contiguous() + } + fn as_contiguous_mut(&self) -> Option> { + self.0.as_contiguous_mut() + } + fn to_contiguous(&self) -> Vec { + self.0.to_contiguous() + } +} pub trait Buffer: Debug + PyThreadingConstraint { fn get_options(&self) -> BorrowedValue; @@ -168,6 +212,29 @@ impl PyMemoryView { }) } + pub fn from_buffer_range( + obj: PyObjectRef, + buffer: BufferRef, + range: std::ops::Range, + vm: &VirtualMachine, + ) -> PyResult { + let options = buffer.get_options().clone(); + let itemsize = options.itemsize; + let format_spec = Self::parse_format(&options.format, vm)?; + Ok(PyMemoryView { + obj, + buffer, + options, + released: AtomicCell::new(false), + start: range.start * itemsize, + stop: range.end * itemsize, + step: 1, + exports: AtomicCell::new(0), + format_spec, + hash: OnceCell::new(), + }) + } + pub fn try_bytes(&self, vm: &VirtualMachine, f: F) -> PyResult where F: FnOnce(&[u8]) -> R,