Skip to content

Commit fd10823

Browse files
authored
Preallocate binary joins and preserve buffer lifetimes (#8880)
Assisted-by: Codex:GPT-6
1 parent 47eb950 commit fd10823

8 files changed

Lines changed: 245 additions & 66 deletions

File tree

‎crates/vm/src/anystr.rs‎

Lines changed: 0 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -321,23 +321,6 @@ pub(crate) trait AnyStr {
321321
self.py_pad(width - len, 0, fillchar)
322322
}
323323

324-
fn py_join(
325-
&self,
326-
mut iter: impl core::iter::Iterator<Item = PyResult<impl AnyStrWrapper<Self> + TryFromObject>>,
327-
) -> PyResult<Self::Container> {
328-
let mut joined = if let Some(elem) = iter.next() {
329-
elem?.as_ref().unwrap().to_container()
330-
} else {
331-
return Ok(Self::Container::new());
332-
};
333-
for elem in iter {
334-
let elem = elem?;
335-
joined.push_str(self);
336-
joined.push_str(elem.as_ref().unwrap());
337-
}
338-
Ok(joined)
339-
}
340-
341324
fn py_partition<'a, F, S>(
342325
&'a self,
343326
sub: &Self,

‎crates/vm/src/builtins/bytearray.rs‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ use crate::{
1212
bytes_inner::{
1313
ByteInnerFindOptions, ByteInnerHexOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
1414
ByteInnerReplaceOptions, ByteInnerSplitOptions, ByteInnerStripOptions, ByteInnerSub,
15-
ByteInnerTranslateOptions, DecodeArgs, PyBytesInner, bytes_decode,
15+
ByteInnerTranslateOptions, BytesJoin, DecodeArgs, PyBytesInner, bytes_decode,
1616
},
1717
class::PyClassImpl,
1818
common::{
@@ -354,10 +354,15 @@ impl PyByteArray {
354354

355355
#[pymethod]
356356
fn join(&self, iterable_of_bytes: PyObjectRef, vm: &VirtualMachine) -> PyResult<Self> {
357-
// Driving the iterable runs Python, which can reach this bytearray,
358-
// so the separator is taken by value rather than left borrowed.
359-
let separator = self.inner().clone();
360-
Ok(separator.join(iterable_of_bytes, vm)?.into())
357+
// Export before driving the iterable: resizing is forbidden, but
358+
// same-size mutations must be reflected in the joined result.
359+
self.exports.fetch_add(1, Ordering::Release);
360+
scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); }
361+
let separator_len = self.__len__();
362+
let items = BytesJoin::new(iterable_of_bytes, vm)?;
363+
Ok(items
364+
.join(separator_len, || self.borrow_buf().into(), vm)?
365+
.into())
361366
}
362367

363368
#[pymethod]

‎crates/vm/src/builtins/bytes.rs‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ use crate::{
1212
bytes_inner::{
1313
ByteInnerFindOptions, ByteInnerHexOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
1414
ByteInnerReplaceOptions, ByteInnerSplitOptions, ByteInnerStripOptions, ByteInnerSub,
15-
ByteInnerTranslateOptions, DecodeArgs, PyBytesInner, bytes_decode,
15+
ByteInnerTranslateOptions, BytesJoin, DecodeArgs, PyBytesInner, bytes_decode,
1616
},
1717
class::{PyClassDef, PyClassImpl},
1818
common::{hash::PyHash, lock::PyMutex},
@@ -350,8 +350,14 @@ impl PyBytes {
350350
}
351351

352352
#[pymethod]
353-
fn join(&self, iterable_of_bytes: PyObjectRef, vm: &VirtualMachine) -> PyResult<Self> {
354-
Ok(self.inner.join(iterable_of_bytes, vm)?.into())
353+
fn join(&self, iterable_of_bytes: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyBytesRef> {
354+
let items = BytesJoin::new(iterable_of_bytes, vm)?;
355+
if let Some(only) = items.single_bytes(vm) {
356+
return Ok(only);
357+
}
358+
Ok(vm
359+
.ctx
360+
.new_bytes(items.join(self.as_bytes().len(), || self.as_bytes().into(), vm)?))
355361
}
356362

357363
#[pymethod]

‎crates/vm/src/builtins/str.rs‎

Lines changed: 2 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1218,25 +1218,8 @@ impl PyStr {
12181218
// and which has its own wording for what it cannot iterate.
12191219
let iterable = ArgIterable::<PyObjectRef>::try_from_object(vm, iterable)
12201220
.map_err(|_| vm.new_type_error("can only join an iterable"))?;
1221-
let iter = iterable.iter_sized(vm)?.enumerate().map(|(i, obj)| {
1222-
obj?.downcast::<Self>().map_err(|obj| {
1223-
vm.new_type_error(format!(
1224-
"sequence item {i}: expected str instance, {} found",
1225-
obj.class().slot_name()
1226-
))
1227-
})
1228-
});
1229-
let joined = match iter.exactly_one() {
1230-
Ok(first) => {
1231-
let first = first?;
1232-
if first.as_object().class().is(vm.ctx.types.str_type) {
1233-
return Ok(first);
1234-
}
1235-
first.as_wtf8().to_owned()
1236-
}
1237-
Err(iter) => zelf.as_wtf8().py_join(iter)?,
1238-
};
1239-
Ok(vm.ctx.new_str(joined))
1221+
let items = iterable.iter_sized(vm)?.collect::<PyResult<Vec<_>>>()?;
1222+
Self::join_items(&zelf, &items, vm)
12401223
}
12411224

12421225
/// `join` over already-materialized items: checks them and sizes the result before copying.

‎crates/vm/src/bytes_inner.rs‎

Lines changed: 119 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,12 @@ use crate::{
44
VirtualMachine,
55
anystr::{self, AnyStr, AnyStrContainer, AnyStrWrapper},
66
builtins::{
7-
PyBaseExceptionRef, PyByteArray, PyBytes, PyBytesRef, PyInt, PyIntRef, PyStr, PyStrRef,
8-
pystr, pystr::PyUtf8StrRef,
7+
PyBaseExceptionRef, PyByteArray, PyBytes, PyBytesRef, PyInt, PyIntRef, PyList, PyListRef,
8+
PyStr, PyStrRef, PyTuple, PyTupleRef, pystr, pystr::PyUtf8StrRef,
99
},
1010
cformat::cformat_bytes,
11-
common::hash,
1211
common::wtf8::is_py_ascii_whitespace,
12+
common::{borrow::BorrowedValue, hash},
1313
function::{ArgIterable, Either, OptionalArg, PyComparisonValue},
1414
literal::escape::Escape,
1515
protocol::{BufferFlags, PyBuffer},
@@ -40,6 +40,122 @@ impl From<Vec<u8>> for PyBytesInner {
4040
}
4141
}
4242

43+
// Keep exact lists live while acquiring buffers: an exporter can replace items
44+
// or resize the list. No list lock may span a call to the buffer protocol.
45+
pub(crate) struct BytesJoin {
46+
items: Either<PyListRef, PyTupleRef>,
47+
}
48+
49+
impl BytesJoin {
50+
pub(crate) fn new(iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult<Self> {
51+
let items = if let Some(list) = iterable.downcast_ref_if_exact::<PyList>(vm) {
52+
Either::A(list.to_owned())
53+
} else if let Some(tuple) = iterable.downcast_ref_if_exact::<PyTuple>(vm) {
54+
Either::B(tuple.to_owned())
55+
} else {
56+
// PySequence_Fast exhausts the iterator before inspecting its items.
57+
let iterable = ArgIterable::<PyObjectRef>::try_from_object(vm, iterable)
58+
.map_err(|_| vm.new_type_error("can only join an iterable"))?;
59+
let items = iterable.iter_sized(vm)?.collect::<PyResult<Vec<_>>>()?;
60+
Either::B(vm.ctx.new_tuple(items))
61+
};
62+
Ok(Self { items })
63+
}
64+
65+
fn len(&self) -> usize {
66+
match &self.items {
67+
Either::A(list) => list.borrow_vec().len(),
68+
Either::B(tuple) => tuple.as_slice().len(),
69+
}
70+
}
71+
72+
fn item(&self, i: usize) -> Option<PyObjectRef> {
73+
match &self.items {
74+
Either::A(list) => list.borrow_vec().get(i).cloned(),
75+
Either::B(tuple) => tuple.as_slice().get(i).cloned(),
76+
}
77+
}
78+
79+
pub(crate) fn single_bytes(&self, vm: &VirtualMachine) -> Option<PyBytesRef> {
80+
if self.len() != 1 {
81+
return None;
82+
}
83+
self.item(0)?
84+
.downcast_ref_if_exact::<PyBytes>(vm)
85+
.map(ToOwned::to_owned)
86+
}
87+
88+
// stringlib_bytes_join
89+
pub(crate) fn join<'a>(
90+
&self,
91+
separator_len: usize,
92+
mut separator: impl FnMut() -> BorrowedValue<'a, [u8]>,
93+
vm: &VirtualMachine,
94+
) -> PyResult<Vec<u8>> {
95+
let count = self.len();
96+
let mut objects = Vec::with_capacity(count);
97+
let mut buffers = Vec::new();
98+
let mut len = 0usize;
99+
let overflow = || vm.new_overflow_error("join() result is too long");
100+
let changed = || vm.new_runtime_error("sequence changed size during iteration");
101+
for i in 0..count {
102+
let obj = self.item(i).ok_or_else(changed)?;
103+
let item_len = if let Some(bytes) = obj.downcast_ref_if_exact::<PyBytes>(vm) {
104+
bytes.as_bytes().len()
105+
} else {
106+
// Whatever the buffer lookup ran into, the item is only ever
107+
// reported as the wrong kind of thing.
108+
let type_error = || {
109+
vm.new_type_error(format!(
110+
"sequence item {i}: expected a bytes-like object, {} found",
111+
obj.class().slot_name()
112+
))
113+
};
114+
let buffer = PyBuffer::from_object(vm, &obj, BufferFlags::SIMPLE)
115+
.map_err(|_| type_error())?;
116+
if !buffer.desc.is_contiguous() {
117+
return Err(type_error());
118+
}
119+
let item_len = buffer.desc.len;
120+
buffers.push((i, buffer));
121+
item_len
122+
};
123+
objects.push(obj);
124+
len = len.checked_add(item_len).ok_or_else(overflow)?;
125+
if i != 0 {
126+
len = len.checked_add(separator_len).ok_or_else(overflow)?;
127+
}
128+
if len > isize::MAX as usize {
129+
return Err(overflow());
130+
}
131+
if self.len() != count {
132+
return Err(changed());
133+
}
134+
}
135+
let mut joined = Vec::new();
136+
joined
137+
.try_reserve_exact(len)
138+
.map_err(|_| vm.new_memory_error(""))?;
139+
// All Python callbacks have finished. Read the separator only now,
140+
// so same-size mutations made by the iterator or exporters are visible.
141+
let mut buffers_iter = buffers.iter().peekable();
142+
for (i, obj) in objects.iter().enumerate() {
143+
if i != 0 {
144+
// Drop this borrow before acquiring an input buffer: it may
145+
// refer to the same bytearray, with a writer waiting on it.
146+
joined.extend_from_slice(&separator());
147+
}
148+
if buffers_iter.peek().is_some_and(|(index, _)| *index == i) {
149+
let (_, buffer) = buffers_iter.next().unwrap();
150+
joined.extend_from_slice(&buffer.as_contiguous().unwrap());
151+
} else {
152+
joined.extend_from_slice(obj.downcast_ref::<PyBytes>().unwrap().as_bytes());
153+
}
154+
}
155+
Ok(joined)
156+
}
157+
}
158+
43159
/// "y*": any bytes-like object, and nothing else.
44160
impl<'a> TryFromBorrowedObject<'a> for PyBytesInner {
45161
fn try_from_borrowed_object(vm: &VirtualMachine, obj: &'a PyObject) -> PyResult<Self> {
@@ -624,25 +740,6 @@ impl PyBytesInner {
624740
.py_count(needle.as_slice(), range, |h, n| h.find_iter(n).count()))
625741
}
626742

627-
// stringlib_bytes_join
628-
pub fn join(&self, iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
629-
// `PySequence_Fast()`, as in `PyUnicode_Join()`.
630-
let iterable = ArgIterable::<PyObjectRef>::try_from_object(vm, iterable)
631-
.map_err(|_| vm.new_type_error("can only join an iterable"))?;
632-
let iter = iterable.iter_sized(vm)?.enumerate().map(|(i, obj)| {
633-
let obj = obj?;
634-
// Whatever the buffer lookup ran into, the item is only ever
635-
// reported as the wrong kind of thing.
636-
Self::try_from_object(vm, obj.clone()).map_err(|_| {
637-
vm.new_type_error(format!(
638-
"sequence item {i}: expected a bytes-like object, {} found",
639-
obj.class().slot_name()
640-
))
641-
})
642-
});
643-
self.elements.py_join(iter)
644-
}
645-
646743
#[inline]
647744
pub fn find<F>(
648745
&self,

‎extra_tests/snippets/builtin_bytearray.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,29 @@ class B(bytearray):
318318
b"123abcxyz"
319319
)
320320

321+
join_separator = bytearray(b",")
322+
323+
324+
def join_mutate_separator():
325+
yield b"a"
326+
join_separator[0] = ord(":")
327+
yield b"b"
328+
329+
330+
assert join_separator.join(join_mutate_separator()) == b"a:b"
331+
332+
333+
def join_resize_separator():
334+
yield b"a"
335+
join_separator.extend(b"!")
336+
yield b"b"
337+
338+
339+
with assert_raises(BufferError):
340+
join_separator.join(join_resize_separator())
341+
join_separator.extend(b"!")
342+
assert join_separator == b":!"
343+
321344

322345
# endswith startswith
323346
assert bytearray(b"abcde").endswith(b"de")

‎extra_tests/snippets/builtin_bytes.py‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -300,6 +300,77 @@
300300
assert b"abc".join((b"123", b"xyz")) == b"123abcxyz"
301301

302302

303+
class JoinBytes(bytes):
304+
def __buffer__(self, flags):
305+
return memoryview(b"override")
306+
307+
308+
class JoinBuffer:
309+
def __init__(self, action=None):
310+
self.action = action
311+
312+
def __buffer__(self, flags):
313+
if self.action is None:
314+
join_chunk[0] = ord("b")
315+
else:
316+
self.action()
317+
return memoryview(b"x")
318+
319+
def __release_buffer__(self, view):
320+
join_events.append("release")
321+
322+
323+
for join_type in (bytes, bytearray):
324+
join_separator = join_type(b",")
325+
326+
def join_mutable_chunks():
327+
chunk = bytearray(b"a")
328+
yield chunk
329+
chunk[:] = b"bc"
330+
yield chunk
331+
332+
assert join_separator.join(join_mutable_chunks()) == b"bc,bc"
333+
334+
def join_broken_iterable():
335+
yield 42
336+
yield b"a"
337+
raise RuntimeError("producer failed")
338+
339+
with assert_raises(RuntimeError):
340+
join_separator.join(join_broken_iterable())
341+
342+
for sequence_type in (list, tuple):
343+
344+
class JoinSequence(sequence_type):
345+
def __iter__(self):
346+
return iter([b"a", b"b"])
347+
348+
assert join_separator.join(JoinSequence([42])) == b"a,b"
349+
350+
result = join_separator.join([JoinBytes(b"original")])
351+
assert result == b"override" and type(result) is join_type
352+
353+
join_chunk = bytearray(b"a")
354+
join_events = []
355+
assert join_separator.join([join_chunk, JoinBuffer()]) == b"b,x"
356+
assert join_events == ["release"]
357+
with assert_raises(TypeError):
358+
join_separator.join([join_chunk, JoinBuffer(lambda: join_chunk.extend(b"c"))])
359+
join_chunk.extend(b"c")
360+
join_items = []
361+
join_exporter = JoinBuffer(join_items.clear)
362+
join_items[:] = [join_exporter, b"y"]
363+
with assert_raises(RuntimeError):
364+
join_separator.join(join_items)
365+
join_events.clear()
366+
with assert_raises(TypeError):
367+
join_separator.join([JoinBuffer(), 42])
368+
assert join_events == ["release"]
369+
370+
join_single = b"single item"
371+
assert b",".join([join_single]) is join_single
372+
373+
303374
# endswith startswith
304375
assert b"abcde".endswith(b"de")
305376
assert b"abcde".endswith(b"")

‎extra_tests/snippets/builtin_str.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -294,6 +294,17 @@ class JoinStr(str):
294294
assert type("-".join((JoinStr("a"), "b"))) is str
295295
single = "single"
296296
assert "-".join([single]) is single
297+
298+
299+
def join_broken_iterable():
300+
yield 42
301+
yield "a"
302+
raise RuntimeError("producer failed")
303+
304+
305+
with assert_raises(RuntimeError):
306+
"-".join(join_broken_iterable())
307+
297308
assert "HALLO".isupper()
298309
assert not "123".isupper()
299310
assert not "123".islower()

0 commit comments

Comments
 (0)