Merge pull request RustPython#3029 from youknowone/pybuffer

Refactor PyBuffer

Author: youknowone

Lib/test/test_httplib.py

@@ -839,8 +839,6 @@ def test_send_type_error(self):
         with self.assertRaises(TypeError):
             conn.request('POST', 'test', conn)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_chunked(self):
         expected = chunked_expected
         sock = FakeSocket(chunked_start + last_chunk + chunked_end)
@@ -873,8 +871,6 @@ def test_chunked(self):
             finally:
                 resp.close()
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_readinto_chunked(self):
 
         expected = chunked_expected

Lib/test/test_io.py

@@ -1039,8 +1039,6 @@ def test_destructor(self):
         super().test_destructor(self)
 
 class PyIOTest(IOTest):
-    # TODO: RUSTPYTHON, can't resize b/c of existing exports
-    @unittest.expectedFailure
     def test_optional_abilities(self):
         super().test_optional_abilities()
 
@@ -2630,8 +2628,7 @@ def test_non_text_encoding_codecs_are_rejected(self):
         with self.assertRaisesRegex(LookupError, "is not a text encoding"):
             self.TextIOWrapper(b, encoding="hex")
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
+    @unittest.skip('TODO: RUSTPYTHON')
     def test_detach(self):
         r = self.BytesIO()
         b = self.BufferedWriter(r)
@@ -2693,8 +2690,7 @@ def test_line_buffering(self):
         t.write("A\rB")
         self.assertEqual(r.getvalue(), b"XY\nZA\rB")
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
+    @unittest.skip('TODO: RUSTPYTHON')
     def test_reconfigure_line_buffering(self):
         r = self.BytesIO()
         b = self.BufferedWriter(r, 1000)
@@ -3924,16 +3920,12 @@ class PyTextIOWrapperTest(TextIOWrapperTest):
     def test_newlines(self):
         super().test_newlines()
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_line_buffering(self):
         super().test_line_buffering()
 
     def test_seeking_too(self):
         super().test_seeking_too()
 
-    # TODO: RUSTPYTHON, can't resize b/c of existing exports
-    @unittest.expectedFailure
     def test_bufio_write_through(self):
         super().test_bufio_write_through()
 

vm/src/buffer.rs

@@ -3,12 +3,10 @@
 use crate::common::borrow::{BorrowedValue, BorrowedValueMut};
 use crate::common::rc::PyRc;
 use crate::PyThreadingConstraint;
-use crate::{PyObjectRef, PyResult, TypeProtocol};
-use crate::{TryFromBorrowedObject, VirtualMachine};
-use std::{borrow::Cow, fmt::Debug, ops::Deref};
+use crate::{PyObjectRef, PyResult, TryFromBorrowedObject, TypeProtocol, VirtualMachine};
+use std::{borrow::Cow, fmt::Debug};
 
-pub trait PyBuffer: Debug + PyThreadingConstraint {
-    fn get_options(&self) -> &BufferOptions;
+pub trait PyBufferInternal: 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.
@@ -18,49 +16,84 @@ pub trait PyBuffer: Debug + PyThreadingConstraint {
     /// 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);
+}
+
+#[derive(Debug)]
+pub struct PyBuffer {
+    pub obj: PyObjectRef,
+    pub options: BufferOptions,
+    pub(crate) internal: PyRc<dyn PyBufferInternal>,
+}
 
-    fn as_contiguous(&self) -> Option<BorrowedValue<[u8]>> {
-        if !self.get_options().contiguous {
+impl PyBuffer {
+    pub fn new(
+        obj: PyObjectRef,
+        buffer: impl PyBufferInternal + 'static,
+        options: BufferOptions,
+    ) -> Self {
+        buffer.retain();
+        Self {
+            obj,
+            options,
+            internal: PyRc::new(buffer),
+        }
+    }
+    pub fn as_contiguous(&self) -> Option<BorrowedValue<[u8]>> {
+        if !self.options.contiguous {
             return None;
         }
-        Some(self.obj_bytes())
+        Some(self.internal.obj_bytes())
     }
 
-    fn as_contiguous_mut(&self) -> Option<BorrowedValueMut<[u8]>> {
-        if !self.get_options().contiguous {
+    pub fn as_contiguous_mut(&self) -> Option<BorrowedValueMut<[u8]>> {
+        if !self.options.contiguous {
             return None;
         }
-        Some(self.obj_bytes_mut())
+        Some(self.internal.obj_bytes_mut())
     }
 
-    fn to_contiguous(&self) -> Vec<u8> {
-        self.obj_bytes().to_vec()
+    pub fn to_contiguous(&self) -> Vec<u8> {
+        self.internal.obj_bytes().to_vec()
+    }
+
+    pub fn clone_with_options(&self, options: BufferOptions) -> Self {
+        self.internal.retain();
+        Self {
+            obj: self.obj.clone(),
+            options,
+            internal: self.internal.clone(),
+        }
     }
 }
 
 #[derive(Debug, Clone)]
 pub struct BufferOptions {
-    pub readonly: bool,
+    // buf
     pub len: usize,
+    pub readonly: bool,
     pub itemsize: usize,
-    pub contiguous: bool,
     pub format: Cow<'static, str>,
-    // TODO: support multiple dimension array
-    pub ndim: usize,
+    pub ndim: usize, // TODO: support multiple dimension array
     pub shape: Vec<usize>,
     pub strides: Vec<isize>,
+    // suboffsets
+
+    // RustPython fields
+    pub contiguous: bool,
 }
 
 impl BufferOptions {
     pub const DEFAULT: Self = BufferOptions {
-        readonly: true,
         len: 0,
+        readonly: true,
         itemsize: 1,
-        contiguous: true,
         format: Cow::Borrowed("B"),
         ndim: 1,
         shape: Vec::new(),
         strides: Vec::new(),
+        contiguous: true,
     };
 }
 
@@ -70,15 +103,12 @@ impl Default for BufferOptions {
     }
 }
 
-#[derive(Debug)]
-pub struct PyBufferRef(Box<dyn PyBuffer>);
-
-impl TryFromBorrowedObject for PyBufferRef {
+impl TryFromBorrowedObject for PyBuffer {
     fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
         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).map(|x| PyBufferRef(x));
+                return f(obj, vm);
             }
         }
         Err(vm.new_type_error(format!(
@@ -88,76 +118,18 @@ impl TryFromBorrowedObject for PyBufferRef {
     }
 }
 
-impl Drop for PyBufferRef {
-    fn drop(&mut self) {
-        self.0.release();
-    }
-}
-
-impl Deref for PyBufferRef {
-    type Target = dyn PyBuffer;
-
-    fn deref(&self) -> &Self::Target {
-        self.0.deref()
-    }
-}
-
-impl PyBufferRef {
-    pub fn new(buffer: impl PyBuffer + 'static) -> Self {
-        Self(Box::new(buffer))
-    }
-
-    pub fn into_rcbuf(self) -> RcBuffer {
-        // move self.0 out of self; PyBufferRef 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 PyBuffer>> for PyBufferRef {
-    fn from(buffer: Box<dyn PyBuffer>) -> Self {
-        PyBufferRef(buffer)
-    }
-}
-
-#[derive(Debug, Clone)]
-pub struct RcBuffer(PyRc<dyn PyBuffer>);
-impl Deref for RcBuffer {
-    type Target = dyn PyBuffer;
-    fn deref(&self) -> &Self::Target {
-        self.0.deref()
-    }
-}
-
-impl Drop for RcBuffer {
+// What we actually want to implement is:
+// impl<T> Drop for T where T: PyBufferInternal
+// but it is not supported by Rust
+impl Drop for PyBuffer {
     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()
-        }
+        self.internal.release();
     }
 }
 
-impl PyBuffer for RcBuffer {
-    fn get_options(&self) -> &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()
+impl Clone for PyBuffer {
+    fn clone(&self) -> Self {
+        self.clone_with_options(self.options.clone())
     }
 }
 

vm/src/builtins/bytearray.rs

@@ -6,7 +6,7 @@ use super::pystr::PyStrRef;
 use super::pytype::PyTypeRef;
 use super::tuple::PyTupleRef;
 use crate::anystr::{self, AnyStr};
-use crate::buffer::{BufferOptions, PyBuffer, ResizeGuard};
+use crate::buffer::{BufferOptions, PyBuffer, PyBufferInternal, ResizeGuard};
 use crate::bytesinner::{
     bytes_decode, bytes_from_object, value_from_object, ByteInnerFindOptions, ByteInnerNewOptions,
     ByteInnerPaddingOptions, ByteInnerSplitOptions, ByteInnerTranslateOptions, DecodeArgs,
@@ -670,41 +670,35 @@ impl Comparable for PyByteArray {
 }
 
 impl AsBuffer for PyByteArray {
-    fn get_buffer(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<Box<dyn PyBuffer>> {
-        zelf.exports.fetch_add(1);
-        let buf = ByteArrayBuffer {
-            bytearray: zelf.clone(),
-            options: BufferOptions {
+    fn get_buffer(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {
+        let buffer = PyBuffer::new(
+            zelf.as_object().clone(),
+            zelf.clone(),
+            BufferOptions {
                 readonly: false,
                 len: zelf.len(),
                 ..Default::default()
             },
-        };
-        Ok(Box::new(buf))
+        );
+        Ok(buffer)
     }
 }
 
-#[derive(Debug)]
-struct ByteArrayBuffer {
-    bytearray: PyByteArrayRef,
-    options: BufferOptions,
-}
-
-impl PyBuffer for ByteArrayBuffer {
+impl PyBufferInternal for PyRef<PyByteArray> {
     fn obj_bytes(&self) -> BorrowedValue<[u8]> {
-        self.bytearray.borrow_buf().into()
+        self.borrow_buf().into()
     }
 
     fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> {
-        PyRwLockWriteGuard::map(self.bytearray.inner_mut(), |inner| &mut *inner.elements).into()
+        PyRwLockWriteGuard::map(self.inner_mut(), |inner| &mut *inner.elements).into()
     }
 
     fn release(&self) {
-        self.bytearray.exports.fetch_sub(1);
+        self.exports.fetch_sub(1);
     }
 
-    fn get_options(&self) -> &BufferOptions {
-        &self.options
+    fn retain(&self) {
+        self.exports.fetch_add(1);
     }
 }