@@ -18,6 +18,7 @@ mod _pickle {
1818 protocol:: { PyBuffer , PyIter , PyIterReturn } ,
1919 types:: { AsBuffer , Constructor , Initializer , Representable } ,
2020 } ;
21+ use alloc:: collections:: TryReserveError ;
2122 use core:: sync:: atomic:: { AtomicI32 , Ordering } ;
2223 use malachite_bigint:: BigInt ;
2324 use num_traits:: { ToPrimitive , Zero } ;
@@ -407,17 +408,26 @@ mod _pickle {
407408 impl UnpicklerMemoProxy {
408409 #[ pymethod]
409410 fn clear ( zelf : & Py < Self > ) {
410- zelf. unpickler . memo . write ( ) . clear ( ) ;
411+ let old_memo = core:: mem:: take ( & mut * zelf. unpickler . memo . write ( ) ) ;
412+ drop ( old_memo) ;
411413 }
412414
413415 #[ pymethod]
414416 fn copy ( zelf : & Py < Self > , vm : & VirtualMachine ) -> PyResult < PyObjectRef > {
417+ let mut snapshot = Vec :: new ( ) ;
418+ let copied = {
419+ let memo = zelf. unpickler . memo . read ( ) ;
420+ snapshot. try_reserve_exact ( memo. len ( ) ) . map ( |( ) | {
421+ snapshot. extend ( memo. iter ( ) . map ( |( & index, obj) | ( index, obj. clone ( ) ) ) ) ;
422+ } )
423+ } ;
424+ copied. map_err ( |_| vm. no_memory_error ( ) ) ?;
425+ // CPython inserts present memo indices in ascending order.
426+ snapshot. sort_unstable_by_key ( |( index, _) | * index) ;
415427 let dict = vm. ctx . new_dict ( ) ;
416- for ( i, entry) in zelf. unpickler . memo . read ( ) . iter ( ) . enumerate ( ) {
417- if let Some ( obj) = entry {
418- let key: PyObjectRef = vm. ctx . new_int ( i) . into ( ) ;
419- dict. set_item ( & * key, obj. clone ( ) , vm) ?;
420- }
428+ for ( i, obj) in snapshot {
429+ let key: PyObjectRef = vm. ctx . new_int ( i) . into ( ) ;
430+ dict. set_item ( & * key, obj, vm) ?;
421431 }
422432 Ok ( dict. into ( ) )
423433 }
@@ -440,6 +450,10 @@ mod _pickle {
440450
441451 // Unpickler
442452
453+ // PUT indices need not be contiguous. A sparse table also keeps MEMOIZE's
454+ // index equal to the number of entries rather than the largest index.
455+ type UnpicklerMemo = HashMap < usize , PyObjectRef > ;
456+
443457 #[ derive( Debug ) ]
444458 pub ( super ) struct UnpicklerConfig {
445459 initialized : bool ,
@@ -468,7 +482,7 @@ mod _pickle {
468482 #[ derive( Debug , PyPayload ) ]
469483 pub ( super ) struct PyUnpickler {
470484 read_state : PyMutex < ReadState > ,
471- memo : PyRwLock < Vec < Option < PyObjectRef > > > ,
485+ memo : PyRwLock < UnpicklerMemo > ,
472486 config : PyRwLock < UnpicklerConfig > ,
473487 }
474488
@@ -495,7 +509,7 @@ mod _pickle {
495509 fn py_new ( _cls : & Py < PyType > , _args : Self :: Args , _vm : & VirtualMachine ) -> PyResult < Self > {
496510 Ok ( Self {
497511 read_state : PyMutex :: new ( ReadState :: default ( ) ) ,
498- memo : PyRwLock :: new ( Vec :: new ( ) ) ,
512+ memo : PyRwLock :: new ( UnpicklerMemo :: new ( ) ) ,
499513 config : PyRwLock :: new ( UnpicklerConfig :: default ( ) ) ,
500514 } )
501515 }
@@ -548,7 +562,7 @@ mod _pickle {
548562 peek,
549563 ..ReadState :: default ( )
550564 } ;
551- zelf. memo . write ( ) . clear ( ) ;
565+ let old_memo = core :: mem :: take ( & mut * zelf. memo . write ( ) ) ;
552566 * zelf. config . write ( ) = UnpicklerConfig {
553567 initialized : true ,
554568 proto : 0 ,
@@ -557,6 +571,8 @@ mod _pickle {
557571 errors,
558572 buffers,
559573 } ;
574+ drop ( state) ;
575+ drop ( old_memo) ;
560576 Ok ( ( ) )
561577 }
562578 }
@@ -605,21 +621,21 @@ mod _pickle {
605621 }
606622 } ;
607623 let new_memo = if let Some ( proxy) = value. downcast_ref :: < UnpicklerMemoProxy > ( ) {
608- proxy. unpickler . memo . read ( ) . clone ( )
624+ let snapshot = copy_memo ( & proxy. unpickler . memo . read ( ) ) ;
625+ snapshot. map_err ( |_| vm. no_memory_error ( ) ) ?
609626 } else if let Some ( dict) = value. downcast_ref :: < PyDict > ( ) {
610- let mut memo: Vec < Option < PyObjectRef > > = Vec :: new ( ) ;
627+ let mut memo = UnpicklerMemo :: new ( ) ;
628+ memo. try_reserve ( dict. __len__ ( ) )
629+ . map_err ( |_| vm. no_memory_error ( ) ) ?;
611630 for ( key, val) in dict {
612631 let idx = key
613632 . downcast_ref :: < PyInt > ( )
614633 . ok_or_else ( || vm. new_type_error ( "memo key must be integers" ) ) ?;
615- let idx = idx
616- . as_bigint ( )
617- . to_usize ( )
618- . ok_or_else ( || vm. new_value_error ( "memo key must be positive integers." ) ) ?;
619- if memo. len ( ) <= idx {
620- memo. resize ( idx + 1 , None ) ;
634+ let idx = idx. try_to_primitive :: < isize > ( vm) ?;
635+ if idx < 0 {
636+ return Err ( vm. new_value_error ( "memo key must be positive integers." ) ) ;
621637 }
622- memo[ idx] = Some ( val) ;
638+ memo_insert ( & mut memo, idx as usize , & val) . map_err ( |_| vm . no_memory_error ( ) ) ? ;
623639 }
624640 memo
625641 } else {
@@ -628,7 +644,8 @@ mod _pickle {
628644 value. class( ) . name( )
629645 ) ) ) ;
630646 } ;
631- * zelf. memo . write ( ) = new_memo;
647+ let old_memo = core:: mem:: replace ( & mut * zelf. memo . write ( ) , new_memo) ;
648+ drop ( old_memo) ;
632649 Ok ( ( ) )
633650 }
634651 }
@@ -1253,13 +1270,13 @@ mod _pickle {
12531270 Some ( BINPUT ) => {
12541271 let i = st. read_n ( 1 , vm) ?[ 0 ] as usize ;
12551272 let value = top ! ( ) ;
1256- memo_put ( zelf, i , value) ;
1273+ memo_put ( zelf, Some ( i ) , value, vm ) ? ;
12571274 }
12581275 Some ( LONG_BINPUT ) => {
12591276 let arr: [ u8 ; 4 ] = st. read_n ( 4 , vm) ?. try_into ( ) . unwrap ( ) ;
12601277 let i = u32:: from_le_bytes ( arr) as usize ;
12611278 let value = top ! ( ) ;
1262- memo_put ( zelf, i , value) ;
1279+ memo_put ( zelf, Some ( i ) , value, vm ) ? ;
12631280 }
12641281 Some ( PUT ) => {
12651282 let line = st. read_line ( vm) ?;
@@ -1269,12 +1286,11 @@ mod _pickle {
12691286 let body = & line[ ..line. len ( ) - 1 ] ;
12701287 let idx = parse_memo_index ( body, vm) ?;
12711288 let value = top ! ( ) ;
1272- memo_put ( zelf, idx, value) ;
1289+ memo_put ( zelf, Some ( idx) , value, vm ) ? ;
12731290 }
12741291 Some ( MEMOIZE ) => {
12751292 let value = top ! ( ) ;
1276- let idx = zelf. memo . read ( ) . len ( ) ;
1277- memo_put ( zelf, idx, value) ;
1293+ memo_put ( zelf, None , value, vm) ?;
12781294 }
12791295 Some ( PERSID ) => {
12801296 let line = st. read_line ( vm) ?;
@@ -1503,19 +1519,43 @@ mod _pickle {
15031519 }
15041520
15051521 fn memo_get ( zelf : & Py < PyUnpickler > , idx : usize , vm : & VirtualMachine ) -> PyResult < PyObjectRef > {
1506- zelf. memo
1507- . read ( )
1508- . get ( idx)
1509- . and_then ( |o| o. clone ( ) )
1522+ let value = zelf. memo . read ( ) . get ( & idx) . cloned ( ) ;
1523+ value
15101524 . ok_or_else ( || new_unpickling_error ( vm, format ! ( "Memo value not found at index {idx}" ) ) )
15111525 }
15121526
1513- fn memo_put ( zelf : & Py < PyUnpickler > , idx : usize , value : PyObjectRef ) {
1514- let mut memo = zelf. memo . write ( ) ;
1515- if memo. len ( ) <= idx {
1516- memo. resize ( idx + 1 , None ) ;
1527+ fn copy_memo ( memo : & UnpicklerMemo ) -> Result < UnpicklerMemo , TryReserveError > {
1528+ let mut copy = UnpicklerMemo :: new ( ) ;
1529+ copy. try_reserve ( memo. len ( ) ) ?;
1530+ copy. extend ( memo. iter ( ) . map ( |( & idx, obj) | ( idx, obj. clone ( ) ) ) ) ;
1531+ Ok ( copy)
1532+ }
1533+
1534+ fn memo_insert (
1535+ memo : & mut UnpicklerMemo ,
1536+ idx : usize ,
1537+ value : & PyObjectRef ,
1538+ ) -> Result < Option < PyObjectRef > , TryReserveError > {
1539+ if !memo. contains_key ( & idx) {
1540+ memo. try_reserve ( 1 ) ?;
15171541 }
1518- memo[ idx] = Some ( value) ;
1542+ Ok ( memo. insert ( idx, value. clone ( ) ) )
1543+ }
1544+
1545+ fn memo_put (
1546+ zelf : & Py < PyUnpickler > ,
1547+ idx : Option < usize > ,
1548+ value : PyObjectRef ,
1549+ vm : & VirtualMachine ,
1550+ ) -> PyResult < ( ) > {
1551+ let old_value = {
1552+ let mut memo = zelf. memo . write ( ) ;
1553+ let idx = idx. unwrap_or_else ( || memo. len ( ) ) ;
1554+ memo_insert ( & mut memo, idx, & value)
1555+ } ;
1556+ // Releasing an entry can call a finalizer that accesses the memo.
1557+ drop ( old_value. map_err ( |_| vm. no_memory_error ( ) ) ?) ;
1558+ Ok ( ( ) )
15191559 }
15201560
15211561 fn instantiate (
@@ -1699,7 +1739,7 @@ mod _pickle {
16991739 buf : data,
17001740 ..ReadState :: default ( )
17011741 } ) ,
1702- memo : PyRwLock :: new ( Vec :: new ( ) ) ,
1742+ memo : PyRwLock :: new ( UnpicklerMemo :: new ( ) ) ,
17031743 config : PyRwLock :: new ( config) ,
17041744 }
17051745 . into_ref ( & vm. ctx ) ;
@@ -1745,7 +1785,7 @@ mod _pickle {
17451785 peek,
17461786 ..ReadState :: default ( )
17471787 } ) ,
1748- memo : PyRwLock :: new ( Vec :: new ( ) ) ,
1788+ memo : PyRwLock :: new ( UnpicklerMemo :: new ( ) ) ,
17491789 config : PyRwLock :: new ( config) ,
17501790 }
17511791 . into_ref ( & vm. ctx ) ;
0 commit comments