Preallocate binary joins and preserve buffer lifetimes (#8880)

Assisted-by: Codex:GPT-6

crates/vm/src/anystr.rs

@@ -321,23 +321,6 @@ pub(crate) trait AnyStr {
         self.py_pad(width - len, 0, fillchar)
     }
 
-    fn py_join(
-        &self,
-        mut iter: impl core::iter::Iterator<Item = PyResult<impl AnyStrWrapper<Self> + TryFromObject>>,
-    ) -> PyResult<Self::Container> {
-        let mut joined = if let Some(elem) = iter.next() {
-            elem?.as_ref().unwrap().to_container()
-        } else {
-            return Ok(Self::Container::new());
-        };
-        for elem in iter {
-            let elem = elem?;
-            joined.push_str(self);
-            joined.push_str(elem.as_ref().unwrap());
-        }
-        Ok(joined)
-    }
-
     fn py_partition<'a, F, S>(
         &'a self,
         sub: &Self,

crates/vm/src/builtins/bytearray.rs

@@ -12,7 +12,7 @@ use crate::{
     bytes_inner::{
         ByteInnerFindOptions, ByteInnerHexOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
         ByteInnerReplaceOptions, ByteInnerSplitOptions, ByteInnerStripOptions, ByteInnerSub,
-        ByteInnerTranslateOptions, DecodeArgs, PyBytesInner, bytes_decode,
+        ByteInnerTranslateOptions, BytesJoin, DecodeArgs, PyBytesInner, bytes_decode,
     },
     class::PyClassImpl,
     common::{
@@ -354,10 +354,15 @@ impl PyByteArray {
 
     #[pymethod]
     fn join(&self, iterable_of_bytes: PyObjectRef, vm: &VirtualMachine) -> PyResult<Self> {
-        // Driving the iterable runs Python, which can reach this bytearray,
-        // so the separator is taken by value rather than left borrowed.
-        let separator = self.inner().clone();
-        Ok(separator.join(iterable_of_bytes, vm)?.into())
+        // Export before driving the iterable: resizing is forbidden, but
+        // same-size mutations must be reflected in the joined result.
+        self.exports.fetch_add(1, Ordering::Release);
+        scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); }
+        let separator_len = self.__len__();
+        let items = BytesJoin::new(iterable_of_bytes, vm)?;
+        Ok(items
+            .join(separator_len, || self.borrow_buf().into(), vm)?
+            .into())
     }
 
     #[pymethod]

crates/vm/src/builtins/bytes.rs

@@ -12,7 +12,7 @@ use crate::{
     bytes_inner::{
         ByteInnerFindOptions, ByteInnerHexOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
         ByteInnerReplaceOptions, ByteInnerSplitOptions, ByteInnerStripOptions, ByteInnerSub,
-        ByteInnerTranslateOptions, DecodeArgs, PyBytesInner, bytes_decode,
+        ByteInnerTranslateOptions, BytesJoin, DecodeArgs, PyBytesInner, bytes_decode,
     },
     class::{PyClassDef, PyClassImpl},
     common::{hash::PyHash, lock::PyMutex},
@@ -350,8 +350,14 @@ impl PyBytes {
     }
 
     #[pymethod]
-    fn join(&self, iterable_of_bytes: PyObjectRef, vm: &VirtualMachine) -> PyResult<Self> {
-        Ok(self.inner.join(iterable_of_bytes, vm)?.into())
+    fn join(&self, iterable_of_bytes: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyBytesRef> {
+        let items = BytesJoin::new(iterable_of_bytes, vm)?;
+        if let Some(only) = items.single_bytes(vm) {
+            return Ok(only);
+        }
+        Ok(vm
+            .ctx
+            .new_bytes(items.join(self.as_bytes().len(), || self.as_bytes().into(), vm)?))
     }
 
     #[pymethod]

crates/vm/src/builtins/str.rs

@@ -1218,25 +1218,8 @@ impl PyStr {
         // and which has its own wording for what it cannot iterate.
         let iterable = ArgIterable::<PyObjectRef>::try_from_object(vm, iterable)
             .map_err(|_| vm.new_type_error("can only join an iterable"))?;
-        let iter = iterable.iter_sized(vm)?.enumerate().map(|(i, obj)| {
-            obj?.downcast::<Self>().map_err(|obj| {
-                vm.new_type_error(format!(
-                    "sequence item {i}: expected str instance, {} found",
-                    obj.class().slot_name()
-                ))
-            })
-        });
-        let joined = match iter.exactly_one() {
-            Ok(first) => {
-                let first = first?;
-                if first.as_object().class().is(vm.ctx.types.str_type) {
-                    return Ok(first);
-                }
-                first.as_wtf8().to_owned()
-            }
-            Err(iter) => zelf.as_wtf8().py_join(iter)?,
-        };
-        Ok(vm.ctx.new_str(joined))
+        let items = iterable.iter_sized(vm)?.collect::<PyResult<Vec<_>>>()?;
+        Self::join_items(&zelf, &items, vm)
     }
 
     /// `join` over already-materialized items: checks them and sizes the result before copying.

crates/vm/src/bytes_inner.rs

@@ -4,12 +4,12 @@ use crate::{
     VirtualMachine,
     anystr::{self, AnyStr, AnyStrContainer, AnyStrWrapper},
     builtins::{
-        PyBaseExceptionRef, PyByteArray, PyBytes, PyBytesRef, PyInt, PyIntRef, PyStr, PyStrRef,
-        pystr, pystr::PyUtf8StrRef,
+        PyBaseExceptionRef, PyByteArray, PyBytes, PyBytesRef, PyInt, PyIntRef, PyList, PyListRef,
+        PyStr, PyStrRef, PyTuple, PyTupleRef, pystr, pystr::PyUtf8StrRef,
     },
     cformat::cformat_bytes,
-    common::hash,
     common::wtf8::is_py_ascii_whitespace,
+    common::{borrow::BorrowedValue, hash},
     function::{ArgIterable, Either, OptionalArg, PyComparisonValue},
     literal::escape::Escape,
     protocol::{BufferFlags, PyBuffer},
@@ -40,6 +40,122 @@ impl From<Vec<u8>> for PyBytesInner {
     }
 }
 
+// Keep exact lists live while acquiring buffers: an exporter can replace items
+// or resize the list. No list lock may span a call to the buffer protocol.
+pub(crate) struct BytesJoin {
+    items: Either<PyListRef, PyTupleRef>,
+}
+
+impl BytesJoin {
+    pub(crate) fn new(iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult<Self> {
+        let items = if let Some(list) = iterable.downcast_ref_if_exact::<PyList>(vm) {
+            Either::A(list.to_owned())
+        } else if let Some(tuple) = iterable.downcast_ref_if_exact::<PyTuple>(vm) {
+            Either::B(tuple.to_owned())
+        } else {
+            // PySequence_Fast exhausts the iterator before inspecting its items.
+            let iterable = ArgIterable::<PyObjectRef>::try_from_object(vm, iterable)
+                .map_err(|_| vm.new_type_error("can only join an iterable"))?;
+            let items = iterable.iter_sized(vm)?.collect::<PyResult<Vec<_>>>()?;
+            Either::B(vm.ctx.new_tuple(items))
+        };
+        Ok(Self { items })
+    }
+
+    fn len(&self) -> usize {
+        match &self.items {
+            Either::A(list) => list.borrow_vec().len(),
+            Either::B(tuple) => tuple.as_slice().len(),
+        }
+    }
+
+    fn item(&self, i: usize) -> Option<PyObjectRef> {
+        match &self.items {
+            Either::A(list) => list.borrow_vec().get(i).cloned(),
+            Either::B(tuple) => tuple.as_slice().get(i).cloned(),
+        }
+    }
+
+    pub(crate) fn single_bytes(&self, vm: &VirtualMachine) -> Option<PyBytesRef> {
+        if self.len() != 1 {
+            return None;
+        }
+        self.item(0)?
+            .downcast_ref_if_exact::<PyBytes>(vm)
+            .map(ToOwned::to_owned)
+    }
+
+    // stringlib_bytes_join
+    pub(crate) fn join<'a>(
+        &self,
+        separator_len: usize,
+        mut separator: impl FnMut() -> BorrowedValue<'a, [u8]>,
+        vm: &VirtualMachine,
+    ) -> PyResult<Vec<u8>> {
+        let count = self.len();
+        let mut objects = Vec::with_capacity(count);
+        let mut buffers = Vec::new();
+        let mut len = 0usize;
+        let overflow = || vm.new_overflow_error("join() result is too long");
+        let changed = || vm.new_runtime_error("sequence changed size during iteration");
+        for i in 0..count {
+            let obj = self.item(i).ok_or_else(changed)?;
+            let item_len = if let Some(bytes) = obj.downcast_ref_if_exact::<PyBytes>(vm) {
+                bytes.as_bytes().len()
+            } else {
+                // Whatever the buffer lookup ran into, the item is only ever
+                // reported as the wrong kind of thing.
+                let type_error = || {
+                    vm.new_type_error(format!(
+                        "sequence item {i}: expected a bytes-like object, {} found",
+                        obj.class().slot_name()
+                    ))
+                };
+                let buffer = PyBuffer::from_object(vm, &obj, BufferFlags::SIMPLE)
+                    .map_err(|_| type_error())?;
+                if !buffer.desc.is_contiguous() {
+                    return Err(type_error());
+                }
+                let item_len = buffer.desc.len;
+                buffers.push((i, buffer));
+                item_len
+            };
+            objects.push(obj);
+            len = len.checked_add(item_len).ok_or_else(overflow)?;
+            if i != 0 {
+                len = len.checked_add(separator_len).ok_or_else(overflow)?;
+            }
+            if len > isize::MAX as usize {
+                return Err(overflow());
+            }
+            if self.len() != count {
+                return Err(changed());
+            }
+        }
+        let mut joined = Vec::new();
+        joined
+            .try_reserve_exact(len)
+            .map_err(|_| vm.new_memory_error(""))?;
+        // All Python callbacks have finished. Read the separator only now,
+        // so same-size mutations made by the iterator or exporters are visible.
+        let mut buffers_iter = buffers.iter().peekable();
+        for (i, obj) in objects.iter().enumerate() {
+            if i != 0 {
+                // Drop this borrow before acquiring an input buffer: it may
+                // refer to the same bytearray, with a writer waiting on it.
+                joined.extend_from_slice(&separator());
+            }
+            if buffers_iter.peek().is_some_and(|(index, _)| *index == i) {
+                let (_, buffer) = buffers_iter.next().unwrap();
+                joined.extend_from_slice(&buffer.as_contiguous().unwrap());
+            } else {
+                joined.extend_from_slice(obj.downcast_ref::<PyBytes>().unwrap().as_bytes());
+            }
+        }
+        Ok(joined)
+    }
+}
+
 /// "y*": any bytes-like object, and nothing else.
 impl<'a> TryFromBorrowedObject<'a> for PyBytesInner {
     fn try_from_borrowed_object(vm: &VirtualMachine, obj: &'a PyObject) -> PyResult<Self> {
@@ -624,25 +740,6 @@ impl PyBytesInner {
             .py_count(needle.as_slice(), range, |h, n| h.find_iter(n).count()))
     }
 
-    // stringlib_bytes_join
-    pub fn join(&self, iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
-        // `PySequence_Fast()`, as in `PyUnicode_Join()`.
-        let iterable = ArgIterable::<PyObjectRef>::try_from_object(vm, iterable)
-            .map_err(|_| vm.new_type_error("can only join an iterable"))?;
-        let iter = iterable.iter_sized(vm)?.enumerate().map(|(i, obj)| {
-            let obj = obj?;
-            // Whatever the buffer lookup ran into, the item is only ever
-            // reported as the wrong kind of thing.
-            Self::try_from_object(vm, obj.clone()).map_err(|_| {
-                vm.new_type_error(format!(
-                    "sequence item {i}: expected a bytes-like object, {} found",
-                    obj.class().slot_name()
-                ))
-            })
-        });
-        self.elements.py_join(iter)
-    }
-
     #[inline]
     pub fn find<F>(
         &self,

extra_tests/snippets/builtin_bytearray.py

@@ -318,6 +318,29 @@ assert bytearray(b"abc").join((bytearray(b"123"), bytearray(b"xyz"))) == bytearr
     b"123abcxyz"
 )
 
+join_separator = bytearray(b",")
+
+
+def join_mutate_separator():
+    yield b"a"
+    join_separator[0] = ord(":")
+    yield b"b"
+
+
+assert join_separator.join(join_mutate_separator()) == b"a:b"
+
+
+def join_resize_separator():
+    yield b"a"
+    join_separator.extend(b"!")
+    yield b"b"
+
+
+with assert_raises(BufferError):
+    join_separator.join(join_resize_separator())
+join_separator.extend(b"!")
+assert join_separator == b":!"
+
 
 # endswith startswith
 assert bytearray(b"abcde").endswith(b"de")

extra_tests/snippets/builtin_bytes.py

@@ -300,6 +300,77 @@ with assert_raises(TypeError):
 assert b"abc".join((b"123", b"xyz")) == b"123abcxyz"
 
 
+class JoinBytes(bytes):
+    def __buffer__(self, flags):
+        return memoryview(b"override")
+
+
+class JoinBuffer:
+    def __init__(self, action=None):
+        self.action = action
+
+    def __buffer__(self, flags):
+        if self.action is None:
+            join_chunk[0] = ord("b")
+        else:
+            self.action()
+        return memoryview(b"x")
+
+    def __release_buffer__(self, view):
+        join_events.append("release")
+
+
+for join_type in (bytes, bytearray):
+    join_separator = join_type(b",")
+
+    def join_mutable_chunks():
+        chunk = bytearray(b"a")
+        yield chunk
+        chunk[:] = b"bc"
+        yield chunk
+
+    assert join_separator.join(join_mutable_chunks()) == b"bc,bc"
+
+    def join_broken_iterable():
+        yield 42
+        yield b"a"
+        raise RuntimeError("producer failed")
+
+    with assert_raises(RuntimeError):
+        join_separator.join(join_broken_iterable())
+
+    for sequence_type in (list, tuple):
+
+        class JoinSequence(sequence_type):
+            def __iter__(self):
+                return iter([b"a", b"b"])
+
+        assert join_separator.join(JoinSequence([42])) == b"a,b"
+
+    result = join_separator.join([JoinBytes(b"original")])
+    assert result == b"override" and type(result) is join_type
+
+    join_chunk = bytearray(b"a")
+    join_events = []
+    assert join_separator.join([join_chunk, JoinBuffer()]) == b"b,x"
+    assert join_events == ["release"]
+    with assert_raises(TypeError):
+        join_separator.join([join_chunk, JoinBuffer(lambda: join_chunk.extend(b"c"))])
+    join_chunk.extend(b"c")
+    join_items = []
+    join_exporter = JoinBuffer(join_items.clear)
+    join_items[:] = [join_exporter, b"y"]
+    with assert_raises(RuntimeError):
+        join_separator.join(join_items)
+    join_events.clear()
+    with assert_raises(TypeError):
+        join_separator.join([JoinBuffer(), 42])
+    assert join_events == ["release"]
+
+join_single = b"single item"
+assert b",".join([join_single]) is join_single
+
+
 # endswith startswith
 assert b"abcde".endswith(b"de")
 assert b"abcde".endswith(b"")

extra_tests/snippets/builtin_str.py

@@ -294,6 +294,17 @@ assert type("-".join([JoinStr("a")])) is str
 assert type("-".join((JoinStr("a"), "b"))) is str
 single = "single"
 assert "-".join([single]) is single
+
+
+def join_broken_iterable():
+    yield 42
+    yield "a"
+    raise RuntimeError("producer failed")
+
+
+with assert_raises(RuntimeError):
+    "-".join(join_broken_iterable())
+
 assert "HALLO".isupper()
 assert not "123".isupper()
 assert not "123".islower()