Optimize bytes-like (l|r)strip (#4500)

dannasman · web-flow · commit a0c34dab547d · 2023-02-17T21:31:12.000+09:00
diff --git a/vm/src/builtins/bytearray.rs b/vm/src/builtins/bytearray.rs
@@ -507,13 +507,37 @@ impl PyByteArray {
     }
 
     #[pymethod]
-    fn lstrip(&self, chars: OptionalOption<PyBytesInner>) -> Self {
-        self.inner().lstrip(chars).into()
+    fn lstrip(
+        zelf: PyRef<Self>,
+        chars: OptionalOption<PyBytesInner>,
+        vm: &VirtualMachine,
+    ) -> PyRef<Self> {
+        let inner = zelf.inner();
+        let stripped = inner.lstrip(chars);
+        let elements = &inner.elements;
+        if stripped == elements {
+            drop(inner);
+            zelf
+        } else {
+            vm.new_pyref(PyByteArray::from(stripped.to_vec()))
+        }
     }
 
     #[pymethod]
-    fn rstrip(&self, chars: OptionalOption<PyBytesInner>) -> Self {
-        self.inner().rstrip(chars).into()
+    fn rstrip(
+        zelf: PyRef<Self>,
+        chars: OptionalOption<PyBytesInner>,
+        vm: &VirtualMachine,
+    ) -> PyRef<Self> {
+        let inner = zelf.inner();
+        let stripped = inner.rstrip(chars);
+        let elements = &inner.elements;
+        if stripped == elements {
+            drop(inner);
+            zelf
+        } else {
+            vm.new_pyref(PyByteArray::from(stripped.to_vec()))
+        }
     }
 
     /// removeprefix($self, prefix, /)
diff --git a/vm/src/builtins/bytes.rs b/vm/src/builtins/bytes.rs
@@ -366,13 +366,31 @@ impl PyBytes {
     }
 
     #[pymethod]
-    fn lstrip(&self, chars: OptionalOption<PyBytesInner>) -> Self {
-        self.inner.lstrip(chars).into()
+    fn lstrip(
+        zelf: PyRef<Self>,
+        chars: OptionalOption<PyBytesInner>,
+        vm: &VirtualMachine,
+    ) -> PyRef<Self> {
+        let stripped = zelf.inner.lstrip(chars);
+        if stripped == zelf.as_bytes() {
+            zelf
+        } else {
+            vm.ctx.new_bytes(stripped.to_vec())
+        }
     }
 
     #[pymethod]
-    fn rstrip(&self, chars: OptionalOption<PyBytesInner>) -> Self {
-        self.inner.rstrip(chars).into()
+    fn rstrip(
+        zelf: PyRef<Self>,
+        chars: OptionalOption<PyBytesInner>,
+        vm: &VirtualMachine,
+    ) -> PyRef<Self> {
+        let stripped = zelf.inner.rstrip(chars);
+        if stripped == zelf.as_bytes().to_vec() {
+            zelf
+        } else {
+            vm.ctx.new_bytes(stripped.to_vec())
+        }
     }
 
     /// removeprefix($self, prefix, /)
diff --git a/vm/src/bytesinner.rs b/vm/src/bytesinner.rs
@@ -576,24 +576,20 @@ impl PyBytesInner {
             .to_vec()
     }
 
-    pub fn lstrip(&self, chars: OptionalOption<PyBytesInner>) -> Vec<u8> {
-        self.elements
-            .py_strip(
-                chars,
-                |s, chars| s.trim_start_with(|c| chars.contains(&(c as u8))),
-                |s| s.trim_start(),
-            )
-            .to_vec()
+    pub fn lstrip(&self, chars: OptionalOption<PyBytesInner>) -> &[u8] {
+        self.elements.py_strip(
+            chars,
+            |s, chars| s.trim_start_with(|c| chars.contains(&(c as u8))),
+            |s| s.trim_start(),
+        )
     }
 
-    pub fn rstrip(&self, chars: OptionalOption<PyBytesInner>) -> Vec<u8> {
-        self.elements
-            .py_strip(
-                chars,
-                |s, chars| s.trim_end_with(|c| chars.contains(&(c as u8))),
-                |s| s.trim_end(),
-            )
-            .to_vec()
+    pub fn rstrip(&self, chars: OptionalOption<PyBytesInner>) -> &[u8] {
+        self.elements.py_strip(
+            chars,
+            |s, chars| s.trim_end_with(|c| chars.contains(&(c as u8))),
+            |s| s.trim_end(),
+        )
     }
 
     // new in Python 3.9
