From c2c5a81e38ec00205d197adcb902a0b8eff10342 Mon Sep 17 00:00:00 2001
From: Kangzhi Shi <shikangzhi@gmail.com>
Date: Thu, 10 Sep 2020 10:15:08 +0200
Subject: [PATCH] Fix bytes stuck on setslice

---
 Lib/test/test_bytes.py     |  1 -
 vm/src/bytesinner.rs       | 22 ++++++++++------------
 vm/src/obj/objbytearray.rs | 17 +++++++++++++----
 3 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/Lib/test/test_bytes.py b/Lib/test/test_bytes.py
index e6f93b78a..381c40d60 100644
--- a/Lib/test/test_bytes.py
+++ b/Lib/test/test_bytes.py
@@ -1371,7 +1371,6 @@ class ByteArrayTest(BaseBytesTest, unittest.TestCase):
                     del b[start:stop:step]
                     self.assertEqual(b, bytearray(L))
 
-    @unittest.skip("TODO: RUSTPYTHON")
     def test_setslice_trap(self):
         # This test verifies that we correctly handle assigning self
         # to a slice of self (the old Lambert Meertens trap).
diff --git a/vm/src/bytesinner.rs b/vm/src/bytesinner.rs
index e7ea69417..d1afd0c04 100644
--- a/vm/src/bytesinner.rs
+++ b/vm/src/bytesinner.rs
@@ -333,7 +333,12 @@ impl PyBytesInner {
         }
     }
 
-    fn setindex(&mut self, int: isize, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
+    pub fn setindex(
+        &mut self,
+        int: isize,
+        object: PyObjectRef,
+        vm: &VirtualMachine,
+    ) -> PyResult<()> {
         if let Some(idx) = self.elements.get_pos(int) {
             let result = match_class!(match object {
                 i @ PyInt => {
@@ -353,7 +358,7 @@ impl PyBytesInner {
         }
     }
 
-    fn setslice(
+    pub fn setslice(
         &mut self,
         slice: PySliceRef,
         object: PyObjectRef,
@@ -379,16 +384,9 @@ impl PyBytesInner {
         self.elements.set_slice_items(vm, &slice, items.as_slice())
     }
 
-    pub fn setitem(
-        &mut self,
-        needle: SequenceIndex,
-        object: PyObjectRef,
-        vm: &VirtualMachine,
-    ) -> PyResult<()> {
-        match needle {
-            SequenceIndex::Int(int) => self.setindex(int, object, vm),
-            SequenceIndex::Slice(slice) => self.setslice(slice, object, vm),
-        }
+    pub fn setslice_from_self(&mut self, slice: PySliceRef, vm: &VirtualMachine) -> PyResult<()> {
+        let items = self.elements.clone();
+        self.elements.set_slice_items(vm, &slice, items.as_slice())
     }
 
     pub fn delitem(&mut self, needle: SequenceIndex, vm: &VirtualMachine) -> PyResult<()> {
diff --git a/vm/src/obj/objbytearray.rs b/vm/src/obj/objbytearray.rs
index 3acce59c3..bd44f7778 100644
--- a/vm/src/obj/objbytearray.rs
+++ b/vm/src/obj/objbytearray.rs
@@ -16,8 +16,8 @@ use crate::bytesinner::{
 use crate::common::cell::{PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard};
 use crate::function::{OptionalArg, OptionalOption};
 use crate::pyobject::{
-    BorrowValue, Either, PyClassImpl, PyComparisonValue, PyContext, PyIterable, PyObjectRef, PyRef,
-    PyResult, PyValue, TryFromObject,
+    BorrowValue, Either, IdProtocol, PyClassImpl, PyComparisonValue, PyContext, PyIterable,
+    PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
 };
 use crate::pystr::{self, PyCommonString};
 use crate::vm::VirtualMachine;
@@ -186,12 +186,21 @@ impl PyByteArray {
 
     #[pymethod(name = "__setitem__")]
     fn setitem(
-        &self,
+        zelf: PyRef<Self>,
         needle: SequenceIndex,
         value: PyObjectRef,
         vm: &VirtualMachine,
     ) -> PyResult<()> {
-        self.borrow_value_mut().setitem(needle, value, vm)
+        match needle {
+            SequenceIndex::Int(int) => zelf.borrow_value_mut().setindex(int, value, vm),
+            SequenceIndex::Slice(slice) => {
+                if zelf.is(&value) {
+                    zelf.borrow_value_mut().setslice_from_self(slice, vm)
+                } else {
+                    zelf.borrow_value_mut().setslice(slice, value, vm)
+                }
+            }
+        }
     }
 
     #[pymethod(name = "__delitem__")]