Add memory::RcBuffer

coolreader18 · coolreader18 · commit 5ddb678acf61 · 2020-11-06T11:11:07.000-06:00
diff --git 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<Box<dyn Buffer>> for BufferRef {
     fn from(buffer: Box<dyn Buffer>) -> Self {
         BufferRef(buffer)
     }
 }
+#[derive(Debug, Clone)]
+pub struct RcBuffer(PyRc<dyn Buffer>);
+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<BufferOptions> {
+        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<BorrowedValue<[u8]>> {
+        self.0.as_contiguous()
+    }
+    fn as_contiguous_mut(&self) -> Option<BorrowedValueMut<[u8]>> {
+        self.0.as_contiguous_mut()
+    }
+    fn to_contiguous(&self) -> Vec<u8> {
+        self.0.to_contiguous()
+    }
+}
 
 pub trait Buffer: Debug + PyThreadingConstraint {
     fn get_options(&self) -> BorrowedValue<BufferOptions>;
@@ -168,6 +212,29 @@ impl PyMemoryView {
         })
     }
 
+    pub fn from_buffer_range(
+        obj: PyObjectRef,
+        buffer: BufferRef,
+        range: std::ops::Range<usize>,
+        vm: &VirtualMachine,
+    ) -> PyResult<Self> {
+        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<F, R>(&self, vm: &VirtualMachine, f: F) -> PyResult<R>
     where
         F: FnOnce(&[u8]) -> R,
