Avoid copying unchanged bytes results (#9149)

Assisted-by: Codex:gpt-6

crates/vm/src/builtins/bytearray.rs

@@ -927,12 +927,12 @@ impl Py<PyByteArray> {
 
     #[pymethod]
     fn strip(&self, options: ByteInnerStripOptions) -> PyByteArray {
-        self.inner().strip(options.bytes).into()
+        self.inner().strip(options.bytes).to_vec().into()
     }
 
     #[pymethod]
     fn removeprefix(&self, prefix: PyBytesInner) -> PyByteArray {
-        self.inner().removeprefix(prefix).into()
+        self.inner().removeprefix(prefix).to_vec().into()
     }
 
     #[pymethod]
@@ -1006,7 +1006,7 @@ impl Py<PyByteArray> {
         options: ByteInnerReplaceOptions,
         vm: &VirtualMachine,
     ) -> PyResult<PyByteArray> {
-        Ok(self.inner().replace(options, vm)?.into())
+        Ok(self.inner().replace(options, vm)?.into_owned().into())
     }
 
     #[pymethod]

crates/vm/src/builtins/bytes.rs

@@ -490,21 +490,6 @@ impl Py<PyBytes> {
         Ok(self.inner.translate(options, vm)?.into())
     }
 
-    #[pymethod]
-    fn strip(&self, options: ByteInnerStripOptions) -> PyBytes {
-        self.inner.strip(options.bytes).into()
-    }
-
-    #[pymethod]
-    fn removeprefix(&self, prefix: PyBytesInner) -> PyBytes {
-        self.inner.removeprefix(prefix).into()
-    }
-
-    #[pymethod]
-    fn removesuffix(&self, suffix: PyBytesInner) -> PyBytes {
-        self.inner.removesuffix(suffix).into()
-    }
-
     #[pymethod]
     fn split(
         &self,
@@ -571,11 +556,6 @@ impl Py<PyBytes> {
         Ok(self.inner.zfill(width, vm)?.into())
     }
 
-    #[pymethod]
-    fn replace(&self, options: ByteInnerReplaceOptions, vm: &VirtualMachine) -> PyResult<PyBytes> {
-        Ok(self.inner.replace(options, vm)?.into())
-    }
-
     #[pymethod]
     fn title(&self) -> PyBytes {
         self.inner.title().into()
@@ -590,6 +570,38 @@ impl Py<PyBytes> {
 
 #[pyclass]
 impl PyRef<PyBytes> {
+    // These operations return a subslice, so equal lengths mean unchanged contents.
+    fn result_from_slice(&self, result: &[u8], vm: &VirtualMachine) -> Self {
+        if result.len() == self.as_bytes().len() && self.class().is(vm.ctx.types.bytes_type) {
+            self.clone()
+        } else {
+            vm.ctx.new_bytes(result.to_vec())
+        }
+    }
+
+    #[pymethod]
+    fn strip(self, options: ByteInnerStripOptions, vm: &VirtualMachine) -> Self {
+        self.result_from_slice(self.inner.strip(options.bytes), vm)
+    }
+
+    #[pymethod]
+    fn removeprefix(self, prefix: PyBytesInner, vm: &VirtualMachine) -> Self {
+        self.result_from_slice(self.inner.removeprefix(prefix), vm)
+    }
+
+    #[pymethod]
+    fn removesuffix(self, suffix: PyBytesInner, vm: &VirtualMachine) -> Self {
+        self.result_from_slice(self.inner.removesuffix(suffix), vm)
+    }
+
+    #[pymethod]
+    fn replace(self, options: ByteInnerReplaceOptions, vm: &VirtualMachine) -> PyResult<Self> {
+        Ok(match self.inner.replace(options, vm)? {
+            alloc::borrow::Cow::Borrowed(result) => self.result_from_slice(result, vm),
+            alloc::borrow::Cow::Owned(result) => vm.ctx.new_bytes(result),
+        })
+    }
+
     #[pymethod]
     fn __bytes__(self, vm: &VirtualMachine) -> Self {
         if self.is(vm.ctx.types.bytes_type) {
@@ -601,22 +613,12 @@ impl PyRef<PyBytes> {
 
     #[pymethod]
     fn lstrip(self, options: ByteInnerStripOptions, vm: &VirtualMachine) -> Self {
-        let stripped = self.payload.inner.lstrip(options.bytes);
-        if stripped == self.as_bytes() {
-            self
-        } else {
-            vm.ctx.new_bytes(stripped.to_vec())
-        }
+        self.result_from_slice(self.inner.lstrip(options.bytes), vm)
     }
 
     #[pymethod]
     fn rstrip(self, options: ByteInnerStripOptions, vm: &VirtualMachine) -> Self {
-        let stripped = self.payload.inner.rstrip(options.bytes);
-        if stripped == self.as_bytes() {
-            self
-        } else {
-            vm.ctx.new_bytes(stripped.to_vec())
-        }
+        self.result_from_slice(self.inner.rstrip(options.bytes), vm)
     }
 
     // currently, only 'utf-8' and 'ascii' implemented

crates/vm/src/bytes_inner.rs

@@ -21,6 +21,7 @@ use crate::{
 /// [`crate::byte::bytearray_from_object`].
 pub(crate) type FromObject = fn(&VirtualMachine, &PyObject) -> PyResult<Vec<u8>>;
 
+use alloc::borrow::Cow;
 use bstr::ByteSlice;
 use itertools::Itertools;
 use malachite_bigint::BigInt;
@@ -834,14 +835,12 @@ impl PyBytesInner {
         res
     }
 
-    pub fn strip(&self, chars: Option<Self>) -> Vec<u8> {
-        self.elements
-            .py_strip(
-                chars,
-                |s, chars| s.trim_with(|c| chars.contains(&(c as u8))),
-                |s| s.trim(),
-            )
-            .to_vec()
+    pub fn strip(&self, chars: Option<Self>) -> &[u8] {
+        self.elements.py_strip(
+            chars,
+            |s, chars| s.trim_with(|c| chars.contains(&(c as u8))),
+            |s| s.trim(),
+        )
     }
 
     pub fn lstrip(&self, chars: Option<Self>) -> &[u8] {
@@ -861,21 +860,19 @@ impl PyBytesInner {
     }
 
     // new in Python 3.9
-    pub fn removeprefix(&self, prefix: Self) -> Vec<u8> {
+    pub fn removeprefix(&self, prefix: Self) -> &[u8] {
         self.elements
             .py_removeprefix(&prefix.elements, prefix.elements.len(), |s, p| {
                 s.starts_with(p)
             })
-            .to_vec()
     }
 
     // new in Python 3.9
-    pub fn removesuffix(&self, suffix: Self) -> Vec<u8> {
+    pub fn removesuffix(&self, suffix: Self) -> &[u8] {
         self.elements
             .py_removesuffix(&suffix.elements, suffix.elements.len(), |s, p| {
                 s.ends_with(p)
             })
-            .to_vec()
     }
 
     pub fn split<F>(
@@ -1000,7 +997,7 @@ impl PyBytesInner {
         result
     }
 
-    fn replace_delete(&self, from: Self, max_count: Option<usize>) -> Vec<u8> {
+    fn replace_delete(&self, from: Self, max_count: Option<usize>) -> Cow<'_, [u8]> {
         let count = count_substring(
             self.elements.as_slice(),
             from.elements.as_slice(),
@@ -1008,7 +1005,7 @@ impl PyBytesInner {
         );
         if count == 0 {
             // no matches
-            return self.elements.clone();
+            return Cow::Borrowed(&self.elements);
         }
 
         let result_len = self.len() - (count * from.len());
@@ -1026,10 +1023,15 @@ impl PyBytesInner {
             }
         }
         result.extend_from_slice(&self.elements[last_end..]);
-        result
+        Cow::Owned(result)
     }
 
-    pub fn replace_in_place(&self, from: Self, to: Self, max_count: Option<usize>) -> Vec<u8> {
+    pub fn replace_in_place(
+        &self,
+        from: Self,
+        to: Self,
+        max_count: Option<usize>,
+    ) -> Cow<'_, [u8]> {
         let len = from.len();
         let mut iter = self.elements.find_iter(&from.elements);
 
@@ -1037,11 +1039,11 @@ impl PyBytesInner {
             let mut new = self.elements.clone();
             new[offset..offset + len].clone_from_slice(to.elements.as_slice());
             if max_count == Some(1) {
-                return new;
+                return Cow::Owned(new);
             }
             new
         } else {
-            return self.elements.clone();
+            return Cow::Borrowed(&self.elements);
         };
 
         let mut count = max_count.unwrap_or(usize::MAX) - 1;
@@ -1052,7 +1054,7 @@ impl PyBytesInner {
                 break;
             }
         }
-        new
+        Cow::Owned(new)
     }
 
     fn replace_general(
@@ -1061,7 +1063,7 @@ impl PyBytesInner {
         to: Self,
         max_count: Option<usize>,
         vm: &VirtualMachine,
-    ) -> PyResult<Vec<u8>> {
+    ) -> PyResult<Cow<'_, [u8]>> {
         let count = count_substring(
             self.elements.as_slice(),
             from.elements.as_slice(),
@@ -1069,7 +1071,7 @@ impl PyBytesInner {
         );
         if count == 0 {
             // no matches, return unchanged
-            return Ok(self.elements.clone());
+            return Ok(Cow::Borrowed(&self.elements));
         }
 
         // Check for overflow
@@ -1097,10 +1099,14 @@ impl PyBytesInner {
             }
         }
         result.extend_from_slice(&self.elements[last_end..]);
-        Ok(result)
+        Ok(Cow::Owned(result))
     }
 
-    pub fn replace(&self, args: ByteInnerReplaceOptions, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
+    pub fn replace(
+        &self,
+        args: ByteInnerReplaceOptions,
+        vm: &VirtualMachine,
+    ) -> PyResult<Cow<'_, [u8]>> {
         let ByteInnerReplaceOptions {
             old: from,
             new: to,
@@ -1110,9 +1116,9 @@ impl PyBytesInner {
         let max_count = if max_count >= 0 {
             if max_count == 0 || (self.elements.is_empty() && !from.is_empty()) {
                 // nothing to do; return the original bytes
-                return Ok(self.elements.clone());
+                return Ok(Cow::Borrowed(&self.elements));
             } else if self.elements.is_empty() && from.is_empty() {
-                return Ok(to.elements);
+                return Ok(Cow::Owned(to.elements));
             }
             Some(max_count as usize)
         } else {
@@ -1123,19 +1129,19 @@ impl PyBytesInner {
         if from.elements.is_empty() {
             if to.elements.is_empty() {
                 // nothing to do; return the original bytes
-                return Ok(self.elements.clone());
+                return Ok(Cow::Borrowed(&self.elements));
             }
             // insert the 'to' bytes everywhere.
             //     >>> b"Python".replace(b"", b".")
             //     b'.P.y.t.h.o.n.'
-            return Ok(self.replace_interleave(to, max_count));
+            return Ok(Cow::Owned(self.replace_interleave(to, max_count)));
         }
 
         // Except for b"".replace(b"", b"A") == b"A" there is no way beyond this
         // point for an empty self bytes to generate a non-empty bytes
         // Special case so the remaining code always gets a non-empty bytes
         if self.elements.is_empty() {
-            return Ok(self.elements.clone());
+            return Ok(Cow::Borrowed(&self.elements));
         }
 
         if to.elements.is_empty() {

extra_tests/snippets/builtin_bytes.py

@@ -767,6 +767,26 @@ class SubBytes(bytes):
     pass
 
 
+# Unchanged operations reuse exact bytes but must not preserve a subclass.
+for method, args in (
+    ("strip", ()),
+    ("lstrip", ()),
+    ("rstrip", ()),
+    ("removeprefix", (b"!",)),
+    ("removesuffix", (b"!",)),
+    ("replace", (b"!", b"?")),
+    ("replace", (b"!", b"")),
+    ("replace", (b"!", b"longer")),
+    ("replace", (b"", b"")),
+    ("replace", (b"a", b"b", 0)),
+):
+    original = b"unchanged payload"
+    assert getattr(original, method)(*args) is original
+    result = getattr(SubBytes(original), method)(*args)
+    assert result == original
+    assert type(result) is bytes
+
+
 b = SubBytes(b"0123abc*&")
 assert id(b) == id(b)
 assert id(b) != id(b * -1)