diff --git a/crates/vm/src/anystr.rs b/crates/vm/src/anystr.rs index cff397965..f1ea7d735 100644 --- a/crates/vm/src/anystr.rs +++ b/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 + TryFromObject>>, - ) -> PyResult { - 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, diff --git a/crates/vm/src/builtins/bytearray.rs b/crates/vm/src/builtins/bytearray.rs index c72cb2689..a2813e9b4 100644 --- a/crates/vm/src/builtins/bytearray.rs +++ b/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 { - // 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] diff --git a/crates/vm/src/builtins/bytes.rs b/crates/vm/src/builtins/bytes.rs index d087ed8d5..86db780fc 100644 --- a/crates/vm/src/builtins/bytes.rs +++ b/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 { - Ok(self.inner.join(iterable_of_bytes, vm)?.into()) + fn join(&self, iterable_of_bytes: PyObjectRef, vm: &VirtualMachine) -> PyResult { + 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] diff --git a/crates/vm/src/builtins/str.rs b/crates/vm/src/builtins/str.rs index cf66f093d..b7bc9f12f 100644 --- a/crates/vm/src/builtins/str.rs +++ b/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::::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::().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::>>()?; + Self::join_items(&zelf, &items, vm) } /// `join` over already-materialized items: checks them and sizes the result before copying. diff --git a/crates/vm/src/bytes_inner.rs b/crates/vm/src/bytes_inner.rs index 392a941a0..f38b4dbd4 100644 --- a/crates/vm/src/bytes_inner.rs +++ b/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> 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, +} + +impl BytesJoin { + pub(crate) fn new(iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult { + let items = if let Some(list) = iterable.downcast_ref_if_exact::(vm) { + Either::A(list.to_owned()) + } else if let Some(tuple) = iterable.downcast_ref_if_exact::(vm) { + Either::B(tuple.to_owned()) + } else { + // PySequence_Fast exhausts the iterator before inspecting its items. + let iterable = ArgIterable::::try_from_object(vm, iterable) + .map_err(|_| vm.new_type_error("can only join an iterable"))?; + let items = iterable.iter_sized(vm)?.collect::>>()?; + 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 { + 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 { + if self.len() != 1 { + return None; + } + self.item(0)? + .downcast_ref_if_exact::(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> { + 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::(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::().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 { @@ -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> { - // `PySequence_Fast()`, as in `PyUnicode_Join()`. - let iterable = ArgIterable::::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( &self, diff --git a/extra_tests/snippets/builtin_bytearray.py b/extra_tests/snippets/builtin_bytearray.py index 9437c23b4..d897ec947 100644 --- a/extra_tests/snippets/builtin_bytearray.py +++ b/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") diff --git a/extra_tests/snippets/builtin_bytes.py b/extra_tests/snippets/builtin_bytes.py index 57694f6d6..f30ae89e0 100644 --- a/extra_tests/snippets/builtin_bytes.py +++ b/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"") diff --git a/extra_tests/snippets/builtin_str.py b/extra_tests/snippets/builtin_str.py index c8801a64c..d114e2d57 100644 --- a/extra_tests/snippets/builtin_str.py +++ 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()