From 4d40a36fd2c8cbbbcc6fa5ebe1aa3a25405f590d Mon Sep 17 00:00:00 2001
From: Kangzhi Shi <shikangzhi@gmail.com>
Date: Sat, 12 Sep 2020 13:38:19 +0200
Subject: [PATCH] Fix byte __imul__

---
 Lib/test/test_bytes.py     |  4 ----
 vm/src/bytesinner.rs       | 14 +-------------
 vm/src/obj/objbytearray.rs |  5 +++--
 3 files changed, 4 insertions(+), 19 deletions(-)

diff --git a/Lib/test/test_bytes.py b/Lib/test/test_bytes.py
index 73ced6174..200a2f6d2 100644
--- a/Lib/test/test_bytes.py
+++ b/Lib/test/test_bytes.py
@@ -1392,8 +1392,6 @@ class ByteArrayTest(BaseBytesTest, unittest.TestCase):
         else:
             self.fail("bytes += unicode didn't raise TypeError")
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_irepeat(self):
         b = bytearray(b"abc")
         b1 = b
@@ -1402,8 +1400,6 @@ class ByteArrayTest(BaseBytesTest, unittest.TestCase):
         self.assertEqual(b, b1)
         self.assertIs(b, b1)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_irepeat_1char(self):
         b = bytearray(b"x")
         b1 = b
diff --git a/vm/src/bytesinner.rs b/vm/src/bytesinner.rs
index 67ba138ad..57a032e45 100644
--- a/vm/src/bytesinner.rs
+++ b/vm/src/bytesinner.rs
@@ -1010,19 +1010,7 @@ impl PyBytesInner {
             // We can multiple an empty vector by any integer, even if it doesn't fit in an isize.
             return;
         }
-
-        if n <= 0 {
-            self.elements.clear();
-        } else {
-            let n = n.to_usize().unwrap(); // always positive by outer if condition
-
-            let old = self.elements.clone();
-
-            self.elements.reserve((n - 1) * old.len());
-            for _ in 1..n {
-                self.elements.extend(&old);
-            }
-        }
+        self.elements = self.repeat(n);
     }
 }
 
diff --git a/vm/src/obj/objbytearray.rs b/vm/src/obj/objbytearray.rs
index bd44f7778..3f5f66ab4 100644
--- a/vm/src/obj/objbytearray.rs
+++ b/vm/src/obj/objbytearray.rs
@@ -574,8 +574,9 @@ impl PyByteArray {
     }
 
     #[pymethod(name = "__imul__")]
-    fn imul(&self, n: isize) {
-        self.borrow_value_mut().irepeat(n)
+    fn imul(zelf: PyRef<Self>, n: isize) -> PyRef<Self> {
+        zelf.borrow_value_mut().irepeat(n);
+        zelf
     }
 
     #[pymethod(name = "__mod__")]