From fdd38b8193553e3bc1eb1a8452057292a88f2d4c Mon Sep 17 00:00:00 2001 From: 1ndahous3 <9205230+1ndahous3@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:31:50 +0300 Subject: [PATCH] Avoid copying unchanged bytes results (#9149) Assisted-by: Codex:gpt-6 --- crates/vm/src/builtins/bytearray.rs | 6 +-- crates/vm/src/builtins/bytes.rs | 66 ++++++++++++++------------- crates/vm/src/bytes_inner.rs | 62 +++++++++++++------------ extra_tests/snippets/builtin_bytes.py | 20 ++++++++ 4 files changed, 91 insertions(+), 63 deletions(-) diff --git a/crates/vm/src/builtins/bytearray.rs b/crates/vm/src/builtins/bytearray.rs index 8ad76aefc..2a47f385c 100644 --- a/crates/vm/src/builtins/bytearray.rs +++ b/crates/vm/src/builtins/bytearray.rs @@ -927,12 +927,12 @@ impl Py { #[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 { options: ByteInnerReplaceOptions, vm: &VirtualMachine, ) -> PyResult { - Ok(self.inner().replace(options, vm)?.into()) + Ok(self.inner().replace(options, vm)?.into_owned().into()) } #[pymethod] diff --git a/crates/vm/src/builtins/bytes.rs b/crates/vm/src/builtins/bytes.rs index f2d06706b..d300996f8 100644 --- a/crates/vm/src/builtins/bytes.rs +++ b/crates/vm/src/builtins/bytes.rs @@ -490,21 +490,6 @@ impl Py { 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 { Ok(self.inner.zfill(width, vm)?.into()) } - #[pymethod] - fn replace(&self, options: ByteInnerReplaceOptions, vm: &VirtualMachine) -> PyResult { - Ok(self.inner.replace(options, vm)?.into()) - } - #[pymethod] fn title(&self) -> PyBytes { self.inner.title().into() @@ -590,6 +570,38 @@ impl Py { #[pyclass] impl PyRef { + // 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 { + 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 { #[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 diff --git a/crates/vm/src/bytes_inner.rs b/crates/vm/src/bytes_inner.rs index 8c52948c8..cc68d5fd0 100644 --- a/crates/vm/src/bytes_inner.rs +++ b/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>; +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) -> Vec { - 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) -> &[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) -> &[u8] { @@ -861,21 +860,19 @@ impl PyBytesInner { } // new in Python 3.9 - pub fn removeprefix(&self, prefix: Self) -> Vec { + 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 { + 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( @@ -1000,7 +997,7 @@ impl PyBytesInner { result } - fn replace_delete(&self, from: Self, max_count: Option) -> Vec { + fn replace_delete(&self, from: Self, max_count: Option) -> 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) -> Vec { + pub fn replace_in_place( + &self, + from: Self, + to: Self, + max_count: Option, + ) -> 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, vm: &VirtualMachine, - ) -> PyResult> { + ) -> PyResult> { 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> { + pub fn replace( + &self, + args: ByteInnerReplaceOptions, + vm: &VirtualMachine, + ) -> PyResult> { 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() { diff --git a/extra_tests/snippets/builtin_bytes.py b/extra_tests/snippets/builtin_bytes.py index f30ae89e0..759f4a54f 100644 --- a/extra_tests/snippets/builtin_bytes.py +++ b/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)