From cccc67bdf9929948976933c3b9a678a346e8bf06 Mon Sep 17 00:00:00 2001 From: youknowdot Date: Sun, 4 Oct 2026 21:25:10 +0900 Subject: [PATCH 1/2] 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 --- crates/stdlib/src/pickle.rs | 110 ++++++++++++------ .../snippets/stdlib_pickle_sparse_memo.py | 11 ++ 2 files changed, 86 insertions(+), 35 deletions(-) create mode 100644 extra_tests/snippets/stdlib_pickle_sparse_memo.py diff --git a/crates/stdlib/src/pickle.rs b/crates/stdlib/src/pickle.rs index 417a5721165..8badf9a85ad 100644 --- a/crates/stdlib/src/pickle.rs +++ b/crates/stdlib/src/pickle.rs @@ -18,6 +18,7 @@ mod _pickle { protocol::{PyBuffer, PyIter, PyIterReturn}, types::{AsBuffer, Constructor, Initializer, Representable}, }; + use alloc::collections::TryReserveError; use core::sync::atomic::{AtomicI32, Ordering}; use malachite_bigint::BigInt; use num_traits::{ToPrimitive, Zero}; @@ -407,17 +408,26 @@ mod _pickle { impl UnpicklerMemoProxy { #[pymethod] fn clear(zelf: &Py) { - zelf.unpickler.memo.write().clear(); + let old_memo = core::mem::take(&mut *zelf.unpickler.memo.write()); + drop(old_memo); } #[pymethod] fn copy(zelf: &Py, vm: &VirtualMachine) -> PyResult { + let mut snapshot = Vec::new(); + let copied = { + let memo = zelf.unpickler.memo.read(); + snapshot.try_reserve_exact(memo.len()).map(|()| { + snapshot.extend(memo.iter().map(|(&index, obj)| (index, obj.clone()))); + }) + }; + copied.map_err(|_| vm.no_memory_error())?; + // CPython inserts present memo indices in ascending order. + snapshot.sort_unstable_by_key(|(index, _)| *index); let dict = vm.ctx.new_dict(); - for (i, entry) in zelf.unpickler.memo.read().iter().enumerate() { - if let Some(obj) = entry { - let key: PyObjectRef = vm.ctx.new_int(i).into(); - dict.set_item(&*key, obj.clone(), vm)?; - } + for (i, obj) in snapshot { + let key: PyObjectRef = vm.ctx.new_int(i).into(); + dict.set_item(&*key, obj, vm)?; } Ok(dict.into()) } @@ -440,6 +450,10 @@ mod _pickle { // Unpickler + // PUT indices need not be contiguous. A sparse table also keeps MEMOIZE's + // index equal to the number of entries rather than the largest index. + type UnpicklerMemo = HashMap; + #[derive(Debug)] pub(super) struct UnpicklerConfig { initialized: bool, @@ -468,7 +482,7 @@ mod _pickle { #[derive(Debug, PyPayload)] pub(super) struct PyUnpickler { read_state: PyMutex, - memo: PyRwLock>>, + memo: PyRwLock, config: PyRwLock, } @@ -495,7 +509,7 @@ mod _pickle { fn py_new(_cls: &Py, _args: Self::Args, _vm: &VirtualMachine) -> PyResult { Ok(Self { read_state: PyMutex::new(ReadState::default()), - memo: PyRwLock::new(Vec::new()), + memo: PyRwLock::new(UnpicklerMemo::new()), config: PyRwLock::new(UnpicklerConfig::default()), }) } @@ -548,7 +562,7 @@ mod _pickle { peek, ..ReadState::default() }; - zelf.memo.write().clear(); + let old_memo = core::mem::take(&mut *zelf.memo.write()); *zelf.config.write() = UnpicklerConfig { initialized: true, proto: 0, @@ -557,6 +571,8 @@ mod _pickle { errors, buffers, }; + drop(state); + drop(old_memo); Ok(()) } } @@ -605,21 +621,21 @@ mod _pickle { } }; let new_memo = if let Some(proxy) = value.downcast_ref::() { - proxy.unpickler.memo.read().clone() + let snapshot = copy_memo(&proxy.unpickler.memo.read()); + snapshot.map_err(|_| vm.no_memory_error())? } else if let Some(dict) = value.downcast_ref::() { - let mut memo: Vec> = Vec::new(); + let mut memo = UnpicklerMemo::new(); + memo.try_reserve(dict.__len__()) + .map_err(|_| vm.no_memory_error())?; for (key, val) in dict { let idx = key .downcast_ref::() .ok_or_else(|| vm.new_type_error("memo key must be integers"))?; - let idx = idx - .as_bigint() - .to_usize() - .ok_or_else(|| vm.new_value_error("memo key must be positive integers."))?; - if memo.len() <= idx { - memo.resize(idx + 1, None); + let idx = idx.try_to_primitive::(vm)?; + if idx < 0 { + return Err(vm.new_value_error("memo key must be positive integers.")); } - memo[idx] = Some(val); + memo_insert(&mut memo, idx as usize, &val).map_err(|_| vm.no_memory_error())?; } memo } else { @@ -628,7 +644,8 @@ mod _pickle { value.class().name() ))); }; - *zelf.memo.write() = new_memo; + let old_memo = core::mem::replace(&mut *zelf.memo.write(), new_memo); + drop(old_memo); Ok(()) } } @@ -1253,13 +1270,13 @@ mod _pickle { Some(BINPUT) => { let i = st.read_n(1, vm)?[0] as usize; let value = top!(); - memo_put(zelf, i, value); + memo_put(zelf, Some(i), value, vm)?; } Some(LONG_BINPUT) => { let arr: [u8; 4] = st.read_n(4, vm)?.try_into().unwrap(); let i = u32::from_le_bytes(arr) as usize; let value = top!(); - memo_put(zelf, i, value); + memo_put(zelf, Some(i), value, vm)?; } Some(PUT) => { let line = st.read_line(vm)?; @@ -1269,12 +1286,11 @@ mod _pickle { let body = &line[..line.len() - 1]; let idx = parse_memo_index(body, vm)?; let value = top!(); - memo_put(zelf, idx, value); + memo_put(zelf, Some(idx), value, vm)?; } Some(MEMOIZE) => { let value = top!(); - let idx = zelf.memo.read().len(); - memo_put(zelf, idx, value); + memo_put(zelf, None, value, vm)?; } Some(PERSID) => { let line = st.read_line(vm)?; @@ -1503,19 +1519,43 @@ mod _pickle { } fn memo_get(zelf: &Py, idx: usize, vm: &VirtualMachine) -> PyResult { - zelf.memo - .read() - .get(idx) - .and_then(|o| o.clone()) + let value = zelf.memo.read().get(&idx).cloned(); + value .ok_or_else(|| new_unpickling_error(vm, format!("Memo value not found at index {idx}"))) } - fn memo_put(zelf: &Py, idx: usize, value: PyObjectRef) { - let mut memo = zelf.memo.write(); - if memo.len() <= idx { - memo.resize(idx + 1, None); + fn copy_memo(memo: &UnpicklerMemo) -> Result { + let mut copy = UnpicklerMemo::new(); + copy.try_reserve(memo.len())?; + copy.extend(memo.iter().map(|(&idx, obj)| (idx, obj.clone()))); + Ok(copy) + } + + fn memo_insert( + memo: &mut UnpicklerMemo, + idx: usize, + value: &PyObjectRef, + ) -> Result, TryReserveError> { + if !memo.contains_key(&idx) { + memo.try_reserve(1)?; } - memo[idx] = Some(value); + Ok(memo.insert(idx, value.clone())) + } + + fn memo_put( + zelf: &Py, + idx: Option, + value: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult<()> { + let old_value = { + let mut memo = zelf.memo.write(); + let idx = idx.unwrap_or_else(|| memo.len()); + memo_insert(&mut memo, idx, &value) + }; + // Releasing an entry can call a finalizer that accesses the memo. + drop(old_value.map_err(|_| vm.no_memory_error())?); + Ok(()) } fn instantiate( @@ -1699,7 +1739,7 @@ mod _pickle { buf: data, ..ReadState::default() }), - memo: PyRwLock::new(Vec::new()), + memo: PyRwLock::new(UnpicklerMemo::new()), config: PyRwLock::new(config), } .into_ref(&vm.ctx); @@ -1745,7 +1785,7 @@ mod _pickle { peek, ..ReadState::default() }), - memo: PyRwLock::new(Vec::new()), + memo: PyRwLock::new(UnpicklerMemo::new()), config: PyRwLock::new(config), } .into_ref(&vm.ctx); diff --git a/extra_tests/snippets/stdlib_pickle_sparse_memo.py b/extra_tests/snippets/stdlib_pickle_sparse_memo.py new file mode 100644 index 00000000000..23d137df43b --- /dev/null +++ b/extra_tests/snippets/stdlib_pickle_sparse_memo.py @@ -0,0 +1,11 @@ +import _pickle +import io + +# BINPUT leaves holes; MEMOIZE uses the number of present entries. +unpickler = _pickle.Unpickler(io.BytesIO(b"\x80\x04]q\x07\x94h\x01\x86.")) +left, right = unpickler.load() +assert left is right +memo = unpickler.memo.copy() +assert list(memo) == [1, 7] +assert memo[1] is left +assert memo[7] is left From f4c7ea16bc3aa55e3dd8b1e67d2082c42a7800eb Mon Sep 17 00:00:00 2001 From: youknowdot Date: Mon, 5 Oct 2026 02:02:12 +0900 Subject: [PATCH 2/2] Move sparse memo regression to standard pickle snippet Use the module-level snippet name requested in review without changing the regression assertions. Assisted-by: Codex:model-version-unavailable --- .../snippets/{stdlib_pickle_sparse_memo.py => stdlib_pickle.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename extra_tests/snippets/{stdlib_pickle_sparse_memo.py => stdlib_pickle.py} (100%) diff --git a/extra_tests/snippets/stdlib_pickle_sparse_memo.py b/extra_tests/snippets/stdlib_pickle.py similarity index 100% rename from extra_tests/snippets/stdlib_pickle_sparse_memo.py rename to extra_tests/snippets/stdlib_pickle.py