Fix buffer protocol and memoryview

qingshi163 · Nov 13, 2021 · Oct 18, 2021 · Oct 19, 2021 · Oct 23, 2021 · Oct 24, 2021
commit 4dccbfb3a634d1e970f59a9b3e20efe43b0e970b
diff --git a/stdlib/src/array.rs b/stdlib/src/array.rs
@@ -1234,24 +1234,21 @@ mod array {
         }
     }
 
-    #[derive(Debug)]
-    struct PyArrayBufferInternal(PyRef<PyArray>);
-
-    impl BufferInternal for PyArrayBufferInternal {
+    impl BufferInternal for PyArray {
         fn obj_bytes(&self) -> BorrowedValue<[u8]> {
-            self.0.get_bytes().into()
+            self.get_bytes().into()
         }
 
         fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> {
-            self.0.get_bytes_mut().into()
+            self.get_bytes_mut().into()
         }
 
         fn release(&self) {
-            self.0.exports.fetch_sub(1);
+            self.exports.fetch_sub(1);
         }
 
         fn retain(&self) {
-            self.0.exports.fetch_add(1);
+            self.exports.fetch_add(1);
         }
     }
 

diff --git a/vm/src/builtins/bytearray.rs b/vm/src/builtins/bytearray.rs
@@ -9,6 +9,7 @@ use crate::common::{
         PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyMutex, PyRwLock, PyRwLockReadGuard,
         PyRwLockWriteGuard,
     },
+    static_cell,
 };
 use crate::{
     anystr::{self, AnyStr},
@@ -20,7 +21,7 @@ use crate::{
     },
     function::{ArgBytesLike, ArgIterable, FuncArgs, IntoPyObject, OptionalArg, OptionalOption},
     protocol::{
-        BufferInternal, BufferOptions, BufferResizeGuard, PyBuffer, PyIterReturn, PyMappingMethods,
+        BufferMethods, BufferOptions, BufferResizeGuard, PyBuffer, PyIterReturn, PyMappingMethods,
     },
     sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex},
     types::{
@@ -33,6 +34,7 @@ use crate::{
 };
 use bstr::ByteSlice;
 use crossbeam_utils::atomic::AtomicCell;
+use rustpython_common::atomic::Radium;
 use std::mem::size_of;
 
 /// "bytearray(iterable_of_ints) -> bytearray\n\
@@ -709,38 +711,39 @@ impl Comparable for PyByteArray {
 }
 
 impl AsBuffer for PyByteArray {
-    fn as_buffer(zelf: &crate::PyObjectView<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {
+    fn as_buffer(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {
+        static_cell! {
+            static METHODS: BufferMethods;
+        }
+        let methods = METHODS.get_or_init(|| BufferMethods {
+            obj_bytes: |zelf| zelf.payload::<Self>().unwrap().borrow_buf().into(),
+            obj_bytes_mut: |zelf| {
+                let zelf = zelf.payload::<Self>().unwrap();
+                PyRwLockWriteGuard::map(zelf.inner_mut(), |x| &mut *x.elements).into()
+            },
+            release: Some(|zelf| {
+                zelf.payload::<Self>().unwrap().exports.fetch_sub(1);
+            }),
+            retain: Some(|zelf| {
+                zelf.payload::<Self>().unwrap().exports.fetch_add(1);
+            }),
+            contiguous: None,
+            contiguous_mut: None,
+            collect_bytes: None,
+        });
         let buffer = PyBuffer::new(
-            zelf.as_object().to_owned(),
-            zelf.to_owned(),
+            zelf.as_object().clone(),
             BufferOptions {
                 readonly: false,
                 len: zelf.len(),
                 ..Default::default()
             },
+            methods,
         );
         Ok(buffer)
     }
 }
 
-impl BufferInternal for PyRef<PyByteArray> {
-    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);
-    }
-
-    fn retain(&self) {
-        self.exports.fetch_add(1);
-    }
-}
-
 impl<'a> BufferResizeGuard<'a> for PyByteArray {
     type Resizable = PyRwLockWriteGuard<'a, PyBytesInner>;
 

diff --git a/vm/src/builtins/bytes.rs b/vm/src/builtins/bytes.rs
@@ -10,7 +10,7 @@ use crate::{
     function::{
         ArgBytesLike, ArgIterable, IntoPyObject, IntoPyResult, OptionalArg, OptionalOption,
     },
-    protocol::{BufferInternal, BufferOptions, PyBuffer, PyIterReturn, PyMappingMethods},
+    protocol::{BufferMethods, BufferOptions, PyBuffer, PyIterReturn, PyMappingMethods},
     types::{
         AsBuffer, AsMapping, Callable, Comparable, Constructor, Hashable, IterNext,
         IterNextIterable, Iterable, PyComparisonOp, Unconstructible,
@@ -541,33 +541,30 @@ impl PyBytes {
     }
 }
 
+static BUFFER_METHODS: BufferMethods = BufferMethods {
+    obj_bytes: |zelf| zelf.payload::<PyBytes>().unwrap().as_bytes().into(),
+    obj_bytes_mut: |_| panic!(),
+    contiguous: None,
+    contiguous_mut: None,
+    collect_bytes: None,
+    release: None,
+    retain: None,
+};
+
 impl AsBuffer for PyBytes {
     fn as_buffer(zelf: &crate::PyObjectView<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {
         let buf = PyBuffer::new(
-            zelf.as_object().to_owned(),
-            zelf.to_owned(),
+            zelf.as_object().clone(),
             BufferOptions {
                 len: zelf.len(),
                 ..Default::default()
             },
+            &BUFFER_METHODS,
         );
         Ok(buf)
     }
 }
 
-impl BufferInternal for PyRef<PyBytes> {
-    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<Self>, _vm: &VirtualMachine) -> PyMappingMethods {
         PyMappingMethods {