Skip to content

Commit cccc67b

Browse files
committed
Use populated entries for unpickler MEMOIZE indexes
A sparse BINPUT followed by MEMOIZE must use the number of occupied memo entries, not the largest index plus one. Store populated entries in a map, preserve ascending memo copies, reserve fallibly, and release displaced references after memo guards end. Keep existing memo assignment/proxy behavior and add one bounded regression for sparse BINPUT, MEMOIZE identity and copy order. No protocol version, GC traversal or unrelated pickle input changes are included. Extracted from the Python 3.15 migration onto the unchanged 3.14 target. Assisted-by: Codex:model-version-unavailable
1 parent f39b054 commit cccc67b

2 files changed

Lines changed: 86 additions & 35 deletions

File tree

‎crates/stdlib/src/pickle.rs‎

Lines changed: 75 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -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);
Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
import _pickle
2+
import io
3+
4+
# BINPUT leaves holes; MEMOIZE uses the number of present entries.
5+
unpickler = _pickle.Unpickler(io.BytesIO(b"\x80\x04]q\x07\x94h\x01\x86."))
6+
left, right = unpickler.load()
7+
assert left is right
8+
memo = unpickler.memo.copy()
9+
assert list(memo) == [1, 7]
10+
assert memo[1] is left
11+
assert memo[7] is left

0 commit comments

Comments
 (0)