From 04d5e12ead697dd25fbdcbe53730a9c667183fea Mon Sep 17 00:00:00 2001 From: 1ndahous3 <9205230+1ndahous3@users.noreply.github.com> Date: Sun, 27 Sep 2026 07:39:29 +0300 Subject: [PATCH] Raise `MemoryError` on impossible allocations and reuse list/tuple buffers (#8840) * Raise MemoryError instead of aborting on impossible allocations Assisted-by: Claude Code:claude-opus-5-5 * use CPython 3.14's integer digit limit for left shifts * size the pow memory check by what malachite allocates * Make sequence allocations fallible and reuse deallocation buffers Assisted-by: Codex:GPT-6 --------- Co-authored-by: Jeong, YunWon <69878+youknowone@users.noreply.github.com> --- crates/vm/src/builtins/bytearray.rs | 8 +- crates/vm/src/builtins/int.rs | 43 ++++++- crates/vm/src/builtins/list.rs | 32 +++-- crates/vm/src/builtins/tuple.rs | 7 +- crates/vm/src/object/core.rs | 33 +++++ crates/vm/src/protocol/sequence.rs | 113 +++++++++++++++++- crates/vm/src/stdlib/_collections.rs | 16 +-- extra_tests/snippets/builtin_bytearray.py | 3 + extra_tests/snippets/builtin_int.py | 9 ++ extra_tests/snippets/builtin_list.py | 37 ++++++ .../snippets/stdlib_collections_deque.py | 4 + 11 files changed, 280 insertions(+), 25 deletions(-) diff --git a/crates/vm/src/builtins/bytearray.rs b/crates/vm/src/builtins/bytearray.rs index f0e94a8e1..c72cb2689 100644 --- a/crates/vm/src/builtins/bytearray.rs +++ b/crates/vm/src/builtins/bytearray.rs @@ -546,7 +546,13 @@ impl PyByteArray { if size < 0 { return Err(vm.new_value_error("bytearray.resize(): new size must be >= 0")); } - self.try_resizable(vm)?.elements.resize(size as usize, 0); + let mut inner = self.try_resizable(vm)?; + let size = size as usize; + let elements = &mut inner.elements; + elements + .try_reserve_exact(size.saturating_sub(elements.len())) + .map_err(|_| vm.no_memory_error())?; + elements.resize(size, 0); Ok(()) } diff --git a/crates/vm/src/builtins/int.rs b/crates/vm/src/builtins/int.rs index dac4e7223..239bd9ec6 100644 --- a/crates/vm/src/builtins/int.rs +++ b/crates/vm/src/builtins/int.rs @@ -152,6 +152,19 @@ fn inner_pow(int1: &BigInt, int2: &BigInt, vm: &VirtualMachine) -> PyResult { float::float_pow(v1, v2, vm) } else { let value = if let Some(v2) = int2.to_u64() { + // malachite builds a power of two at its exact size. Otherwise it allocates + // `bits * v2` bits for the result plus a scratch buffer of up to the same size. + let base_bits = int1.bits(); + if base_bits > 1 { + let bits = if int1.trailing_zeros() == Some(base_bits - 1) { + (base_bits - 1) + .checked_mul(v2) + .and_then(|b| b.checked_add(1)) + } else { + base_bits.checked_mul(v2).and_then(|b| b.checked_mul(2)) + }; + reserve_result_bits(bits, vm)?; + } return Ok(vm.ctx.new_int(Pow::pow(int1, v2)).into()); } else if int1.is_one() { 1 @@ -198,13 +211,39 @@ fn inner_lshift(base: &BigInt, bits: &BigInt, vm: &VirtualMachine) -> PyResult { bits, |base, bits| base << bits, |bits, vm| { - bits.to_usize() - .ok_or_else(|| vm.new_overflow_error("the number is too large to convert to int")) + // CPython's limit: `(i64::MAX - 1) / 30` digits of 30 bits. + const MAX_DIGITS: u128 = (i64::MAX as u128 - 1) / 30; + let digits = bits.to_u128().map(|shift| { + u128::from(base.bits().div_ceil(30)) + shift / 30 + u128::from(shift % 30 != 0) + }); + if digits.is_none_or(|digits| digits > MAX_DIGITS) { + return Err(vm.new_overflow_error("too many digits in integer")); + } + let shift = bits.to_u64().ok_or_else(|| vm.no_memory_error())?; + reserve_result_bits(base.bits().checked_add(shift), vm)?; + usize::try_from(shift).map_err(|_| vm.no_memory_error()) }, vm, ) } +/// Reject sizes that cannot be allocated before calling Malachite's infallible arithmetic. +/// This is only a preflight check: Malachite allocates its own result and scratch buffers, +/// so an allocation failure during the operation can still abort the process. +fn reserve_result_bits(bits: Option, vm: &VirtualMachine) -> PyResult<()> { + // Results below this size are not worth the extra allocation. + const CHECK_FROM_BITS: u64 = 1 << 26; + let bits = bits.ok_or_else(|| vm.no_memory_error())?; + if bits < CHECK_FROM_BITS { + return Ok(()); + } + let words = + usize::try_from(bits.div_ceil(u64::BITS.into())).map_err(|_| vm.no_memory_error())?; + Vec::::new() + .try_reserve_exact(words) + .map_err(|_| vm.no_memory_error()) +} + fn inner_rshift(base: &BigInt, bits: &BigInt, vm: &VirtualMachine) -> PyResult { inner_shift( base, diff --git a/crates/vm/src/builtins/list.rs b/crates/vm/src/builtins/list.rs index 2e478e17f..cc0736abf 100644 --- a/crates/vm/src/builtins/list.rs +++ b/crates/vm/src/builtins/list.rs @@ -73,7 +73,12 @@ unsafe impl Traverse for PyList { // This is safe because during GC collection, the object is unreachable // and no other code should be accessing it. if let Some(mut guard) = self.elements.try_write() { - out.extend(guard.drain(..)); + if out.is_empty() { + // Reuse the allocation so deallocation does not need more memory. + *out = core::mem::take(&mut *guard); + } else { + out.extend(guard.drain(..)); + } } } } @@ -457,10 +462,20 @@ where F: FnMut(PyObjectRef) -> PyResult, { use crate::builtins::PyTuple; + let mut v = Vec::new(); if let Some(tuple) = obj.downcast_ref_if_exact::(vm) { - tuple.as_slice().iter().map(|x| f(x.clone())).collect() + v.try_reserve_exact(tuple.len()) + .map_err(|_| vm.no_memory_error())?; + for x in tuple.as_slice() { + v.push(f(x.clone())?); + } } else if let Some(list) = obj.downcast_ref_if_exact::(vm) { - list.borrow_vec().iter().map(|x| f(x.clone())).collect() + let elements = list.borrow_vec(); + v.try_reserve_exact(elements.len()) + .map_err(|_| vm.no_memory_error())?; + for x in elements.iter() { + v.push(f(x.clone())?); + } } else { let iter = obj.to_owned().get_iter(vm)?; let iter = iter.iter::(vm)?; @@ -469,13 +484,16 @@ where .length_opt(vm) .transpose()? .unwrap_or(0); - let mut v = Vec::with_capacity(len); + v.try_reserve_exact(len).map_err(|_| vm.no_memory_error())?; for x in iter { - v.push(f(x?)?); + let item = f(x?)?; + if v.len() == v.capacity() { + v.try_reserve(1).map_err(|_| vm.no_memory_error())?; + } + v.push(item); } - v.shrink_to_fit(); - Ok(v) } + Ok(v) } impl MutObjectSequenceOp for PyList { diff --git a/crates/vm/src/builtins/tuple.rs b/crates/vm/src/builtins/tuple.rs index f44920ce6..2297e5aaf 100644 --- a/crates/vm/src/builtins/tuple.rs +++ b/crates/vm/src/builtins/tuple.rs @@ -96,7 +96,12 @@ unsafe impl Traverse for PyTuple { fn clear(&mut self, out: &mut Vec) { let elements = core::mem::take(self.elements.get_mut()); - out.extend(elements.into_vec()); + if out.is_empty() { + // Reuse the allocation so deallocation does not need more memory. + *out = elements.into_vec(); + } else { + out.extend(elements.into_vec()); + } } } diff --git a/crates/vm/src/object/core.rs b/crates/vm/src/object/core.rs index 2b5f9cab1..ef3eaa12e 100644 --- a/crates/vm/src/object/core.rs +++ b/crates/vm/src/object/core.rs @@ -3238,6 +3238,39 @@ pub(crate) fn init_type_hierarchy() -> BootstrapTypeHierarchy { mod tests { use super::*; + #[test] + fn clear_reuses_storage_and_preserves_existing_edges() { + use crate::builtins::PyList; + + crate::Interpreter::without_stdlib(Default::default()).enter(|vm| { + for tuple in [false, true] { + for existing in [false, true] { + let elements = vec![vm.ctx.none(), vm.ctx.none()]; + let allocation = elements.as_ptr(); + let mut sequence: Box = if tuple { + Box::new(PyTuple::new_unchecked(elements.into_boxed_slice())) + } else { + Box::new(PyList::from(elements)) + }; + let mut out = if existing { + vec![vm.ctx.new_int(1).into()] + } else { + Vec::new() + }; + sequence.clear(&mut out); + if existing { + assert_eq!(out[0].try_to_value::(vm).unwrap(), 1); + } else { + assert_eq!(out.as_ptr(), allocation); + } + assert_eq!(out.len(), 2 + usize::from(existing)); + assert!(out[usize::from(existing)..].iter().all(|x| vm.is_none(x))); + sequence.traverse(&mut |_| panic!("cleared sequence still owns an edge")); + } + } + }); + } + #[test] fn miri_test_type_initialization() { let hierarchy = init_type_hierarchy(); diff --git a/crates/vm/src/protocol/sequence.rs b/crates/vm/src/protocol/sequence.rs index 32a1c47ed..2065f65df 100644 --- a/crates/vm/src/protocol/sequence.rs +++ b/crates/vm/src/protocol/sequence.rs @@ -429,21 +429,34 @@ impl PySequence<'_> { where F: FnMut(&PyObject) -> PyResult, { + let mut v = Vec::new(); if let Some(tuple) = self.obj.downcast_ref_if_exact::(vm) { - tuple.as_slice().iter().map(|x| f(x.as_ref())).collect() + v.try_reserve_exact(tuple.len()) + .map_err(|_| vm.no_memory_error())?; + for x in tuple.as_slice() { + v.push(f(x.as_ref())?); + } } else if let Some(list) = self.obj.downcast_ref_if_exact::(vm) { - list.borrow_vec().iter().map(|x| f(x.as_ref())).collect() + let elements = list.borrow_vec(); + v.try_reserve_exact(elements.len()) + .map_err(|_| vm.no_memory_error())?; + for x in elements.iter() { + v.push(f(x.as_ref())?); + } } else { let iter = self.obj.to_owned().get_iter(vm)?; let iter = iter.iter::(vm)?; let len = self.length(vm).unwrap_or(0); - let mut v = Vec::with_capacity(len); + v.try_reserve_exact(len).map_err(|_| vm.no_memory_error())?; for x in iter { - v.push(f(x?.as_ref())?); + let item = f(x?.as_ref())?; + if v.len() == v.capacity() { + v.try_reserve(1).map_err(|_| vm.no_memory_error())?; + } + v.push(item); } - v.shrink_to_fit(); - Ok(v) } + Ok(v) } pub fn contains(self, target: &PyObject, vm: &VirtualMachine) -> PyResult { @@ -475,3 +488,91 @@ impl PySequence<'_> { Ok(false) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::{Interpreter, builtins::PyRange}; + use core::cell::Cell; + + #[derive(Debug)] + struct Converted<'a>(&'a Cell); + + impl Drop for Converted<'_> { + fn drop(&mut self) { + self.0.set(self.0.get() + 1); + } + } + + #[test] + fn conversion_and_partial_result_cleanup() { + Interpreter::without_stdlib(Default::default()).enter(|vm| { + let elements: Vec = + (0..20).map(|value| vm.ctx.new_int(value).into()).collect(); + let list = vm.ctx.new_list(elements.clone()); + let iterator = list.as_object().get_iter(vm).unwrap(); + assert!(iterator.sequence_unchecked().length_opt(vm).is_none()); + let sequences: [PyObjectRef; 3] = [ + vm.ctx.new_tuple(elements).into(), + list.clone().into(), + iterator.into(), + ]; + for sequence in sequences { + let values: Vec = sequence + .sequence_unchecked() + .extract(|item| item.try_to_value(vm), vm) + .unwrap(); + assert_eq!(values, (0..20).collect::>()); + } + + let iterator = list.as_object().get_iter(vm).unwrap(); + let error = vm.new_value_error("conversion failed"); + let dropped = Cell::new(0); + let mut calls = 0; + let raised = iterator + .sequence_unchecked() + .extract( + |_| { + calls += 1; + if calls == 3 { + Err(error.clone()) + } else { + Ok(Converted(&dropped)) + } + }, + vm, + ) + .unwrap_err(); + assert!(raised.is(&error)); + assert_eq!(calls, 3); + assert_eq!(dropped.get(), 2); + }); + } + + #[test] + fn capacity_overflow_precedes_conversion() { + Interpreter::without_stdlib(Default::default()).enter(|vm| { + let range = PyRange { + start: vm.ctx.new_int(0), + stop: vm.ctx.new_int(isize::MAX), + step: vm.ctx.new_int(1), + } + .into_ref(&vm.ctx); + let called = Cell::new(false); + // The byte capacity exceeds isize::MAX; no large allocation is attempted. + let error = range + .as_object() + .sequence_unchecked() + .extract( + |_| { + called.set(true); + Ok(0u64) + }, + vm, + ) + .unwrap_err(); + assert!(error.fast_isinstance(vm.ctx.exceptions.memory_error)); + assert!(!called.get()); + }); + } +} diff --git a/crates/vm/src/stdlib/_collections.rs b/crates/vm/src/stdlib/_collections.rs index c0b0fd180..afc91a8f4 100644 --- a/crates/vm/src/stdlib/_collections.rs +++ b/crates/vm/src/stdlib/_collections.rs @@ -396,14 +396,14 @@ mod _collections { if n > 1 && result_len.saturating_mul(size_of::()) >= MAX_MEMORY_SIZE { return Err(vm.no_memory_error()); } - let iter = deque.iter().cycle().take(mul_len); - let skipped = self - .maxlen - .and_then(|maxlen| mul_len.checked_sub(maxlen)) - .unwrap_or(0); - - let deque = iter.skip(skipped).cloned().collect(); - Ok(deque) + // A maxlen keeps only the last `result_len` items; start the cycle where they begin. + let start = (mul_len - result_len).checked_rem(deque.len()).unwrap_or(0); + let mut result = VecDeque::new(); + result + .try_reserve_exact(result_len) + .map_err(|_| vm.no_memory_error())?; + result.extend(deque.iter().cycle().skip(start).take(result_len).cloned()); + Ok(result) } fn __mul__(&self, n: isize, vm: &VirtualMachine) -> PyResult { diff --git a/extra_tests/snippets/builtin_bytearray.py b/extra_tests/snippets/builtin_bytearray.py index 610c9017f..9437c23b4 100644 --- a/extra_tests/snippets/builtin_bytearray.py +++ b/extra_tests/snippets/builtin_bytearray.py @@ -3,6 +3,9 @@ import sys from testutils import assert_raises +with assert_raises(MemoryError): + bytearray(b"ab").resize(sys.maxsize) + # new assert bytearray([1, 2, 3]) assert bytearray((1, 2, 3)) diff --git a/extra_tests/snippets/builtin_int.py b/extra_tests/snippets/builtin_int.py index c111d2532..d887c9a79 100644 --- a/extra_tests/snippets/builtin_int.py +++ b/extra_tests/snippets/builtin_int.py @@ -407,3 +407,12 @@ finally: for value in (0, 1, -1): with assert_raises(MemoryError): value.to_bytes(2**60, "big", signed=True) + +# Results too large to allocate raise MemoryError instead of aborting +assert_raises(MemoryError, lambda: 2**62 << 2**62) +assert_raises(MemoryError, lambda: 1 << 9223372036854775770) +assert_raises(OverflowError, lambda: 1 << 9223372036854775771) +assert_raises(OverflowError, lambda: 1 << 2**64) +assert_raises(OverflowError, lambda: -5 << 10**30) +assert 0 << 2**64 == 0 +assert pow(1, 2**63) == 1 and pow(-1, 2**63 + 1) == -1 and pow(0, 2**63) == 0 diff --git a/extra_tests/snippets/builtin_list.py b/extra_tests/snippets/builtin_list.py index 3cc6eab14..e782bf1f2 100644 --- a/extra_tests/snippets/builtin_list.py +++ b/extra_tests/snippets/builtin_list.py @@ -462,6 +462,43 @@ assert x == [0, 1, 2, 3, 4, "a", "b", "c"] assert x == y assert x == z assert x == zz + + +# A huge length hint is a MemoryError, not an aborted process +class HugeLen: + def __len__(self): + return 2**62 + + def __getitem__(self, i): + if i >= 3: + raise IndexError + return i + + +slice_target = [3, 1, 2] +with assert_raises(MemoryError): + slice_target[1::-1] = HugeLen() + + +class ShortLen(HugeLen): + def __len__(self): + return 1 + + +slice_target[::-1] = ShortLen() +assert slice_target == [2, 1, 0] + + +def fails_during_iteration(): + yield 9 + raise MemoryError + + +with assert_raises(MemoryError): + slice_target[:] = fails_during_iteration() +assert slice_target == [2, 1, 0] + + # insert sec x = a[:] y = a[:] diff --git a/extra_tests/snippets/stdlib_collections_deque.py b/extra_tests/snippets/stdlib_collections_deque.py index 660469eda..aef8f764b 100644 --- a/extra_tests/snippets/stdlib_collections_deque.py +++ b/extra_tests/snippets/stdlib_collections_deque.py @@ -152,6 +152,10 @@ assert list(d) == ["b", "c", "d"] d.appendleft("z") assert list(d) == ["z", "b", "c"] +# Repeating a bounded deque keeps only the last maxlen items without walking the rest +assert deque([1], maxlen=3) * 2**62 == deque([1, 1, 1]) +assert deque([1, 2, 3], maxlen=4) * 3 == deque([3, 1, 2, 3]) + def assert_deque_error(error_type, message, function, *args): with assert_raises(error_type) as caught: