diff --git a/crates/stdlib/src/array.rs b/crates/stdlib/src/array.rs index b51bc02d3fb..6b29d448463 100644 --- a/crates/stdlib/src/array.rs +++ b/crates/stdlib/src/array.rs @@ -75,6 +75,16 @@ mod array { macro_rules! def_array_enum { ($(($n:ident, $t:ty, $c:literal, $scode:literal)),*$(,)?) => { + #[derive(Copy, Clone, Debug, Eq, PartialEq)] + enum ArrayType { + $($n,)* + } + + #[derive(Debug)] + enum ArrayElementValue { + $($n($t),)* + } + #[derive(Debug, Clone)] pub enum ArrayContentType { $($n(Vec<$t>),)* @@ -127,6 +137,12 @@ mod array { } } + fn kind(&self) -> ArrayType { + match self { + $(ArrayContentType::$n(_) => ArrayType::$n,)* + } + } + fn reserve(&mut self, len: usize) { match self { $(ArrayContentType::$n(v) => v.reserve(len),)* @@ -143,6 +159,25 @@ mod array { Ok(()) } + fn truncate(&mut self, len: usize) { + match self { + $(ArrayContentType::$n(v) => v.truncate(len),)* + } + } + + fn set_or_push(&mut self, index: usize, value: ArrayElementValue) { + match (self, value) { + $((ArrayContentType::$n(v), ArrayElementValue::$n(val)) => { + if index < v.len() { + v[index] = val; + } else { + v.push(val); + } + },)* + _ => unreachable!(), + } + } + fn pop(&mut self, i: isize, vm: &VirtualMachine) -> PyResult { match self { $(ArrayContentType::$n(v) => { @@ -483,6 +518,15 @@ mod array { })* } } + + } + + impl ArrayType { + fn convert(self, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult { + match self { + $(ArrayType::$n => <$t>::try_into_from_object(vm, obj).map(ArrayElementValue::$n),)* + } + } } }; } @@ -739,9 +783,22 @@ mod array { self.read().itemsize() } + /// Append a new item to the end of the array. + /// Matches `ins1()` in arraymodule.c. #[pymethod] fn append(zelf: &Py, x: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { - zelf.try_resizable(vm)?.push(x, vm) + let (n, kind) = { + let r = zelf.read(); + (r.len(), r.kind()) + }; + // First conversion (validation, result discarded) + kind.convert(x.clone(), vm)?; + // Truncate to n+1 if array grew during first conversion + zelf.try_resizable(vm)?.truncate(n + 1); + // Second conversion and store at position n + let value = kind.convert(x, vm)?; + zelf.try_resizable(vm)?.set_or_push(n, value); + Ok(()) } #[pymethod] diff --git a/extra_tests/snippets/stdlib_array.py b/extra_tests/snippets/stdlib_array.py index ed2a8f22369..99c7cf4a164 100644 --- a/extra_tests/snippets/stdlib_array.py +++ b/extra_tests/snippets/stdlib_array.py @@ -143,3 +143,27 @@ def write(self, chunk): arr = array("b", range(128)) arr.tofile(_ReenteringWriter(arr)) assert len(arr) == 129 + + +# Regression: `fromfile` should not deadlock when __index__ re-enters append. +class _ReenteringIndex: + def __init__(self, arr): + self.arr = arr + + def __index__(self): + self.arr.append(0) + return 0 + + +class _ReenteringReader: + def __init__(self, arr): + self.arr = arr + + def read(self, n): + self.arr.append(_ReenteringIndex(self.arr)) + return b"\0" * n + + +arr = array("b") +arr.fromfile(_ReenteringReader(arr), 1) +assert list(arr) == [0, 0, 0]