From d1f21793971fb7117fe08f532197613fb264a951 Mon Sep 17 00:00:00 2001
From: Kangzhi Shi <shikangzhi@gmail.com>
Date: Thu, 17 Sep 2020 20:06:08 +0200
Subject: [PATCH] Fix byte pickle implement

---
 extra_tests/snippets/bytearray.py | 13 +++++++++++++
 vm/src/obj/objbytearray.rs        | 17 ++++++-----------
 vm/src/obj/objbytes.rs            | 12 +++---------
 3 files changed, 22 insertions(+), 20 deletions(-)

diff --git a/extra_tests/snippets/bytearray.py b/extra_tests/snippets/bytearray.py
index 31c2ae68b..210da6e3d 100644
--- a/extra_tests/snippets/bytearray.py
+++ b/extra_tests/snippets/bytearray.py
@@ -1,4 +1,5 @@
 from testutils import assert_raises
+import pickle
 
 # new
 assert bytearray([1, 2, 3])
@@ -720,3 +721,15 @@ b = a.copy()
 assert a.__ne__(b) is False
 b = bytearray(b'my bytearray')
 assert a.__ne__(b) is True
+
+# pickle
+a = bytearray(b'\xffab\x80\0\0\370\0\0')
+# TODO: Different output when pickle dumping a str if protocol is 0
+# CPython's output is like:
+# b'c__builtin__\nbytearray\np0\n(c_codecs\nencode\np1\n(V\xffab\x80\\u0000\\u0000\xf8\\u0000\\u0000\np2\nVlatin1\np3\ntp4\nRp5\ntp6\nRp7\n.'
+# We do have fully Interoperability between our output and CPython's
+assert pickle.dumps(a, 0) == b'c__builtin__\nbytearray\np0\n(c_codecs\nencode\np1\n(V\xffab\x80\x00\x00\xf8\x00\x00\np2\nVlatin1\np3\ntp4\nRp5\ntp6\nRp7\n.'
+assert pickle.dumps(a, 1) == b'c__builtin__\nbytearray\nq\x00(c_codecs\nencode\nq\x01(X\x0c\x00\x00\x00\xc3\xbfab\xc2\x80\x00\x00\xc3\xb8\x00\x00q\x02X\x06\x00\x00\x00latin1q\x03tq\x04Rq\x05tq\x06Rq\x07.'
+assert pickle.dumps(a, 2) == b'\x80\x02c__builtin__\nbytearray\nq\x00c_codecs\nencode\nq\x01X\x0c\x00\x00\x00\xc3\xbfab\xc2\x80\x00\x00\xc3\xb8\x00\x00q\x02X\x06\x00\x00\x00latin1q\x03\x86q\x04Rq\x05\x85q\x06Rq\x07.'
+assert pickle.dumps(a, 3) == b'\x80\x03cbuiltins\nbytearray\nq\x00C\t\xffab\x80\x00\x00\xf8\x00\x00q\x01\x85q\x02Rq\x03.'
+assert pickle.dumps(a, 4) == b'\x80\x04\x95*\x00\x00\x00\x00\x00\x00\x00\x8c\x08builtins\x94\x8c\tbytearray\x94\x93\x94C\t\xffab\x80\x00\x00\xf8\x00\x00\x94\x85\x94R\x94.'
diff --git a/vm/src/obj/objbytearray.rs b/vm/src/obj/objbytearray.rs
index b132dfe00..1b68a99d9 100644
--- a/vm/src/obj/objbytearray.rs
+++ b/vm/src/obj/objbytearray.rs
@@ -15,6 +15,7 @@ use crate::bytesinner::{
 use crate::byteslike::PyBytesLike;
 use crate::common::cell::{PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard};
 use crate::function::{OptionalArg, OptionalOption};
+use crate::obj::objbytes::PyBytes;
 use crate::obj::objtuple::PyTuple;
 use crate::pyobject::{
     BorrowValue, Either, IdProtocol, IntoPyObject, PyClassImpl, PyComparisonValue, PyContext,
@@ -566,20 +567,14 @@ impl PyByteArray {
     }
 
     #[pymethod(magic)]
-    fn reduce_ex(&self, _proto: usize, vm: &VirtualMachine) -> (PyClassRef, PyTuple) {
-        self.reduce(vm)
+    fn reduce_ex(zelf: PyRef<Self>, _proto: usize, vm: &VirtualMachine) -> (PyClassRef, PyTuple) {
+        Self::reduce(zelf, vm)
     }
 
     #[pymethod(magic)]
-    fn reduce(&self, vm: &VirtualMachine) -> (PyClassRef, PyTuple) {
-        let param: Vec<PyObjectRef> = self
-            .borrow_value()
-            .elements
-            .iter()
-            .map(|x| x.into_pyobject(vm))
-            .collect();
-        let param = PyTuple::from(vec![vm.ctx.new_tuple(param)]);
-        (Self::class(vm), param)
+    fn reduce(zelf: PyRef<Self>, vm: &VirtualMachine) -> (PyClassRef, PyTuple) {
+        let bytes = PyBytes::from(zelf.borrow_value().elements.clone()).into_pyobject(vm);
+        (Self::class(vm), PyTuple::from(vec![bytes]))
     }
 }
 
diff --git a/vm/src/obj/objbytes.rs b/vm/src/obj/objbytes.rs
index f2f6f409b..99f35be40 100644
--- a/vm/src/obj/objbytes.rs
+++ b/vm/src/obj/objbytes.rs
@@ -14,7 +14,7 @@ use crate::bytesinner::{
 };
 use crate::byteslike::PyBytesLike;
 use crate::function::{OptionalArg, OptionalOption};
-use crate::obj::objtuple::PyTuple;
+use crate::obj::objtuple::{PyTuple, PyTupleRef};
 use crate::pyobject::{
     BorrowValue, Either, IntoPyObject, PyClassImpl, PyComparisonValue, PyContext, PyIterable,
     PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
@@ -466,20 +466,14 @@ impl PyBytes {
     }
 
     #[pymethod(magic)]
-    fn reduce_ex(&self, _proto: usize, vm: &VirtualMachine) -> (PyClassRef, PyTuple) {
-        self.reduce(vm)
-    }
-
-    #[pymethod(magic)]
-    fn reduce(&self, vm: &VirtualMachine) -> (PyClassRef, PyTuple) {
+    fn getnewargs(&self, vm: &VirtualMachine) -> PyTupleRef {
         let param: Vec<PyObjectRef> = self
             .inner
             .elements
             .iter()
             .map(|x| x.into_pyobject(vm))
             .collect();
-        let param = PyTuple::from(vec![vm.ctx.new_tuple(param)]);
-        (Self::class(vm), param)
+        PyTuple::from(param).into_ref(vm)
     }
 }