@@ -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.
44160impl < ' 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 ,
0 commit comments