From 5d227f7ca38e68182390fefb80983f80e5310b55 Mon Sep 17 00:00:00 2001 From: dot Date: Wed, 4 Nov 2026 19:52:42 +0900 Subject: [PATCH 1/6] Align builtin iterator controls and callback lifetimes Share map and zip argument and strict controls, preserve filter subclasses, and release callable iterator state around reentrant callbacks. Assisted-by: OpenAI Codex:model version not exposed --- crates/vm/src/builtins/filter.rs | 35 +++---- crates/vm/src/builtins/iter.rs | 83 ++++++++++----- crates/vm/src/builtins/map.rs | 127 ++++++++++++----------- crates/vm/src/builtins/zip.rs | 168 +++++++++++++++++++++++-------- 4 files changed, 266 insertions(+), 147 deletions(-) diff --git a/crates/vm/src/builtins/filter.rs b/crates/vm/src/builtins/filter.rs index 11e11fc2795..35fb9330c04 100644 --- a/crates/vm/src/builtins/filter.rs +++ b/crates/vm/src/builtins/filter.rs @@ -1,6 +1,6 @@ use super::{PyType, PyTypeRef}; use crate::{ - Context, Py, PyObjectRef, PyPayload, PyResult, VirtualMachine, + AsObject, Context, Py, PyObjectRef, PyPayload, PyResult, VirtualMachine, class::PyClassImpl, protocol::{PyIter, PyIterReturn}, raise_if_stop, @@ -26,7 +26,7 @@ pub struct FilterArgs { #[pyarg(positional)] function: PyObjectRef, #[pyarg(positional)] - iterable: PyIter, + iterable: PyObjectRef, } impl Constructor for PyFilter { @@ -35,25 +35,25 @@ impl Constructor for PyFilter { fn py_new( _cls: &Py, - Self::Args { - function, - iterable: iterator, - }: Self::Args, - _vm: &VirtualMachine, + Self::Args { function, iterable }: Self::Args, + vm: &VirtualMachine, ) -> PyResult { - Ok(Self { - predicate: function, + let result = iterable.get_iter(vm).map(|iterator| Self { + predicate: function.clone(), iterator, - }) + }); + drop(iterable); + drop(function); + result } } #[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE))] impl Py { #[pymethod] - fn __reduce__(&self, vm: &VirtualMachine) -> (PyTypeRef, (PyObjectRef, PyIter)) { + fn __reduce__(&self) -> (PyTypeRef, (PyObjectRef, PyIter)) { ( - vm.ctx.types.filter_type.to_owned(), + self.class().to_owned(), (self.predicate.clone(), self.iterator.clone()), ) } @@ -66,17 +66,18 @@ impl IterNext for PyFilter { let predicate = &zelf.predicate; loop { let next_obj = raise_if_stop!(zelf.iterator.next(vm)?); - let predicate_value = if vm.is_none(predicate) { - next_obj.clone() + let is_true = if vm.is_none(predicate) || predicate.is(vm.ctx.types.bool_type) { + next_obj.try_to_bool(vm)? } else { // the predicate itself can raise StopIteration which does stop the filter iteration - raise_if_stop!(PyIterReturn::from_pyresult( + let predicate_value = raise_if_stop!(PyIterReturn::from_pyresult( predicate.call((next_obj.clone(),), vm), vm - )?) + )?); + predicate_value.try_to_bool(vm)? }; - if predicate_value.try_to_bool(vm)? { + if is_true { return Ok(PyIterReturn::Return(next_obj)); } } diff --git a/crates/vm/src/builtins/iter.rs b/crates/vm/src/builtins/iter.rs index c3eca4dc7f2..365aed9da6f 100644 --- a/crates/vm/src/builtins/iter.rs +++ b/crates/vm/src/builtins/iter.rs @@ -4,14 +4,14 @@ use super::{PyInt, PyTupleRef, PyType}; use crate::{ - Context, Py, PyObject, PyObjectRef, PyPayload, PyResult, VirtualMachine, + AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyResult, VirtualMachine, class::PyClassImpl, function::ArgCallable, object::{Traverse, TraverseFn}, protocol::PyIterReturn, - types::{IterNext, Iterable, SelfIter}, + types::{IterNext, Iterable, PyComparisonOp, SelfIter}, }; -use rustpython_common::lock::{PyMutex, PyRwLock, PyRwLockUpgradableReadGuard}; +use rustpython_common::lock::{PyMutex, PyRwLock}; /// Marks status of iterator. #[derive(Debug, Clone)] @@ -302,8 +302,7 @@ impl IterNext for PySequenceIterator { #[pyclass(module = false, name = "callable_iterator", traverse)] #[derive(Debug)] pub struct PyCallableIterator { - sentinel: PyObjectRef, - status: PyRwLock>, + status: PyRwLock>, } impl PyPayload for PyCallableIterator { @@ -317,10 +316,19 @@ impl PyCallableIterator { #[must_use] pub const fn new(callable: ArgCallable, sentinel: PyObjectRef) -> Self { Self { - sentinel, - status: PyRwLock::new(IterStatus::Active(callable)), + status: PyRwLock::new(IterStatus::Active((callable, sentinel))), } } + + fn exhaust(&self) { + let released = { + let mut status = self.status.write(); + core::mem::replace(&mut *status, IterStatus::Exhausted) + }; + // Dropping the callable or sentinel can invoke a finalizer that calls + // next() again, so both must be released after the state lock. + drop(released); + } } #[pyclass(with(IterNext, Iterable))] @@ -328,10 +336,19 @@ impl Py { #[pymethod] fn __reduce__(&self, vm: &VirtualMachine) -> PyResult { let func = builtins_iter(vm)?; - let status = self.status.read(); - if let IterStatus::Active(callable) = &*status { - let callable_obj: PyObjectRef = callable.clone().into(); - Ok(vm.new_tuple((func, (callable_obj, self.sentinel.clone())))) + let active = { + let status = self.status.read(); + match &*status { + IterStatus::Active((callable, sentinel)) => { + let callable: PyObjectRef = callable.clone().into(); + Some((callable, sentinel.clone())) + } + IterStatus::Exhausted => None, + } + }; + // Tuple allocation may collect objects whose finalizers re-enter next(). + if let Some(args) = active { + Ok(vm.new_tuple((func, args))) } else { Ok(vm.new_tuple((func, (vm.ctx.empty_tuple.clone(),)))) } @@ -346,30 +363,42 @@ impl IterNext for PyCallableIterator { let callable = { let status = zelf.status.read(); match &*status { - IterStatus::Active(callable) => callable.clone(), + IterStatus::Active((callable, _)) => callable.clone(), IterStatus::Exhausted => return Ok(PyIterReturn::StopIteration(None)), } }; - let ret = callable.invoke((), vm)?; - - // Re-check before comparing, but don't hold the lock while running - // sentinel equality. User __eq__ code can re-enter this iterator. - { - let status = zelf.status.read(); - if !matches!(&*status, IterStatus::Active(_)) { + let ret = callable.invoke((), vm); + drop(callable); + let ret = match PyIterReturn::from_pyresult(ret, vm)? { + PyIterReturn::Return(ret) => ret, + PyIterReturn::StopIteration(value) => { + drop(value); + zelf.exhaust(); return Ok(PyIterReturn::StopIteration(None)); } - } - - let is_sentinel = vm.identical_or_equal(&ret, &zelf.sentinel)?; + }; - if is_sentinel { - let status = zelf.status.upgradable_read(); - if !matches!(&*status, IterStatus::Active(_)) { - return Ok(PyIterReturn::StopIteration(None)); + // Re-check after calling, and release the lock before equality: either + // callback can re-enter and exhaust this iterator. + let sentinel = { + let status = zelf.status.read(); + match &*status { + IterStatus::Active((_, sentinel)) => sentinel.clone(), + IterStatus::Exhausted => return Ok(PyIterReturn::StopIteration(None)), } - *PyRwLockUpgradableReadGuard::upgrade(status) = IterStatus::Exhausted; + }; + let is_sentinel = if sentinel.is(&ret) { + drop(sentinel); + true + } else { + // Consume the snapshot before truth conversion: __bool__ can + // exhaust us and observe the sentinel's finalizer immediately. + let comparison = sentinel.rich_compare(ret.clone(), PyComparisonOp::Eq, vm)?; + comparison.try_to_bool(vm)? + }; + if is_sentinel { + zelf.exhaust(); Ok(PyIterReturn::StopIteration(None)) } else { Ok(PyIterReturn::Return(ret)) diff --git a/crates/vm/src/builtins/map.rs b/crates/vm/src/builtins/map.rs index 0a62f2d792b..26aa0a53e48 100644 --- a/crates/vm/src/builtins/map.rs +++ b/crates/vm/src/builtins/map.rs @@ -1,9 +1,12 @@ -use super::{PyType, PyTypeRef}; +use super::{ + PyType, PyTypeRef, + zip::{get_iterators, next_item, release_iterables, strict_value}, +}; use crate::{ - AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, VirtualMachine, + AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, builtins::PyTupleRef, class::PyClassImpl, - function::{ArgIntoBool, FuncArgs, PosArgs}, + function::{FuncArgs, OptionalArg, PosArgs}, protocol::{PyIter, PyIterReturn}, types::{Constructor, IterNext, Iterable, SelfIter}, }; @@ -30,21 +33,21 @@ pub struct PyMapNewArgs { #[pyarg(positional)] function: PyObjectRef, #[pyarg(positional)] - iterable: PyIter, + iterable: PyObjectRef, #[pyarg(flatten)] - iterables: PosArgs, - #[pyarg(named, default)] - strict: bool, + iterables: PosArgs, + #[pyarg(named, default, py_default = "False")] + strict: OptionalArg, } #[derive(FromArgs)] struct MapCallArgs { - #[pyarg(positional)] - function: PyObjectRef, + #[pyarg(positional, optional)] + function: OptionalArg, #[pyarg(flatten)] - iterables: PosArgs, - #[pyarg(named, default)] - strict: bool, + iterables: PosArgs, + #[pyarg(named, default, py_default = "False")] + strict: OptionalArg, } impl Constructor for PyMap { @@ -52,31 +55,61 @@ impl Constructor for PyMap { fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult { let MapCallArgs { - function: mapper, + function, iterables, strict, } = args.bind_for(vm, "map")?; - let iterators = iterables.into_vec(); - if iterators.is_empty() { - return Err(vm.new_type_error("map() must have at least two arguments.")); - } - let payload = Self { - mapper, - iterators, - strict: Radium::new(strict), - }; - payload.into_ref_with_type(vm, cls).map(Into::into) + let result = (|| { + let strict = strict_value(&strict, vm)?; + let mapper = function + .as_ref() + .into_option() + .filter(|_| !iterables.as_ref().is_empty()) + .ok_or_else(|| vm.new_type_error("map() must have at least two arguments."))?; + let payload = Self::from_iterables(mapper, iterables.iter(), strict, vm)?; + payload.into_ref_with_type(vm, cls).map(Into::into) + })(); + release_iterables(strict, iterables); + drop(function); + result } fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { - let PyMapNewArgs { + let Self::Args { function, iterable, iterables, strict, } = args; - let _ = (function, iterable, iterables, strict); - Err(vm.new_type_error("use slot_new")) + let result = (|| { + let strict = strict_value(&strict, vm)?; + Self::from_iterables( + &function, + core::iter::once(&iterable).chain(iterables.iter()), + strict, + vm, + ) + })(); + release_iterables(strict, iterables); + drop(iterable); + drop(function); + result + } +} + +impl PyMap { + fn from_iterables<'a>( + mapper: &PyObjectRef, + iterables: impl IntoIterator, + strict: bool, + vm: &VirtualMachine, + ) -> PyResult { + let iterators = get_iterators(iterables, vm)?; + Ok(Self { + mapper: mapper.clone(), + iterators, + strict: Radium::new(strict), + }) } } @@ -96,10 +129,10 @@ impl PyMap { } #[pymethod] - fn __setstate__(zelf: PyRef, object: PyObjectRef, vm: &VirtualMachine) { - if let Ok(obj) = ArgIntoBool::try_from_object(vm, object) { - zelf.strict.store(obj.into(), atomic::Ordering::Release); - } + fn __setstate__(zelf: PyRef, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + let strict = object.try_to_bool(vm)?; + zelf.strict.store(strict, atomic::Ordering::Release); + Ok(()) } } @@ -107,36 +140,10 @@ impl SelfIter for PyMap {} impl IterNext for PyMap { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { - let mut next_objs = Vec::new(); - for (idx, iterator) in zelf.iterators.iter().enumerate() { - let item = match iterator.next(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => { - if zelf.strict.load(atomic::Ordering::Acquire) { - if idx > 0 { - let plural = if idx == 1 { " " } else { "s 1-" }; - return Err(vm.new_value_error(format!( - "map() argument {} is shorter than argument{}{}", - idx + 1, - plural, - idx, - ))); - } - for (idx, iterator) in zelf.iterators[1..].iter().enumerate() { - if let PyIterReturn::Return(_) = iterator.next(vm)? { - let plural = if idx == 0 { " " } else { "s 1-" }; - return Err(vm.new_value_error(format!( - "map() argument {} is longer than argument{}{}", - idx + 2, - plural, - idx + 1, - ))); - } - } - } - return Ok(PyIterReturn::StopIteration(v)); - } - }; + let mut next_objs = Vec::with_capacity(zelf.iterators.len()); + for idx in 0..zelf.iterators.len() { + let item = + crate::raise_if_stop!(next_item(&zelf.iterators, idx, &zelf.strict, "map", vm)?); next_objs.push(item); } diff --git a/crates/vm/src/builtins/zip.rs b/crates/vm/src/builtins/zip.rs index 3bdaca7f707..e0c900a9d01 100644 --- a/crates/vm/src/builtins/zip.rs +++ b/crates/vm/src/builtins/zip.rs @@ -1,9 +1,9 @@ use super::PyType; use crate::{ - AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, VirtualMachine, + AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, builtins::PyTupleRef, class::PyClassImpl, - function::{ArgIntoBool, PosArgs}, + function::{OptionalArg, PosArgs}, protocol::{PyIter, PyIterReturn}, types::{Constructor, IterNext, Iterable, SelfIter}, }; @@ -27,9 +27,9 @@ impl PyPayload for PyZip { #[derive(FromArgs)] pub struct PyZipNewArgs { #[pyarg(flatten)] - iterables: PosArgs, - #[pyarg(named, default)] - strict: bool, + iterables: PosArgs, + #[pyarg(named, default, py_default = "False")] + strict: OptionalArg, } impl Constructor for PyZip { @@ -37,12 +37,19 @@ impl Constructor for PyZip { fn py_new( _cls: &Py, - PyZipNewArgs { iterables, strict }: Self::Args, - _vm: &VirtualMachine, + Self::Args { iterables, strict }: Self::Args, + vm: &VirtualMachine, ) -> PyResult { - let iterators = iterables.into_vec(); - let strict = Radium::new(strict); - Ok(Self { iterators, strict }) + let result = (|| { + let strict = strict_value(&strict, vm)?; + let iterators = get_iterators(iterables.iter(), vm)?; + Ok(Self { + iterators, + strict: Radium::new(strict), + }) + })(); + release_iterables(strict, iterables); + result } } @@ -65,10 +72,10 @@ impl PyZip { } #[pymethod] - fn __setstate__(zelf: PyRef, object: PyObjectRef, vm: &VirtualMachine) { - if let Ok(obj) = ArgIntoBool::try_from_object(vm, object) { - zelf.strict.store(obj.into(), atomic::Ordering::Release); - } + fn __setstate__(zelf: PyRef, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + let strict = object.try_to_bool(vm)?; + zelf.strict.store(strict, atomic::Ordering::Release); + Ok(()) } } @@ -80,41 +87,116 @@ impl IterNext for PyZip { return Ok(PyIterReturn::StopIteration(None)); } let mut next_objs = Vec::with_capacity(zelf.iterators.len()); - for (idx, iterator) in zelf.iterators.iter().enumerate() { - let item = match iterator.next(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => { - if zelf.strict.load(atomic::Ordering::Acquire) { - if idx > 0 { - let plural = if idx == 1 { " " } else { "s 1-" }; - return Err(vm.new_value_error(format!( - "zip() argument {} is shorter than argument{}{}", - idx + 1, - plural, - idx - ))); - } - for (idx, iterator) in zelf.iterators[1..].iter().enumerate() { - if let PyIterReturn::Return(_obj) = iterator.next(vm)? { - let plural = if idx == 0 { " " } else { "s 1-" }; - return Err(vm.new_value_error(format!( - "zip() argument {} is longer than argument{}{}", - idx + 2, - plural, - idx + 1 - ))); - } - } - } - return Ok(PyIterReturn::StopIteration(v)); - } - }; + for idx in 0..zelf.iterators.len() { + let item = + crate::raise_if_stop!(next_item(&zelf.iterators, idx, &zelf.strict, "zip", vm)?); next_objs.push(item); } Ok(PyIterReturn::Return(vm.ctx.new_tuple(next_objs).into())) } } +pub(super) fn strict_value( + strict: &OptionalArg, + vm: &VirtualMachine, +) -> PyResult { + strict + .as_ref() + .map(|value| value.try_to_bool(vm)) + .transpose() + .map(|value| value.unwrap_or(false)) +} + +pub(super) fn get_iterators<'a>( + iterables: impl IntoIterator, + vm: &VirtualMachine, +) -> PyResult> { + let iterables = iterables.into_iter(); + let mut iterators = Vec::with_capacity(iterables.size_hint().0); + for obj in iterables { + match obj.get_iter(vm) { + Ok(iterator) => iterators.push(iterator), + Err(error) => { + while let Some(iterator) = iterators.pop() { + drop(iterator); + } + return Err(error); + } + } + } + Ok(iterators) +} + +// Constructor callbacks borrow the complete argument list. Release it only +// afterwards, in call-stack order, including when conversion failed. +pub(super) fn release_iterables( + strict: OptionalArg, + iterables: PosArgs, +) { + drop(strict); + iterables.into_iter().rev().for_each(drop); +} + +// Read strict after the source callback, but before clearing a StopIteration +// exception: either operation can run user code that changes the flag. +pub(super) fn next_item( + iterators: &[PyIter], + index: usize, + strict: &PyAtomic, + name: &str, + vm: &VirtualMachine, +) -> PyResult { + let result = iterators[index].next(vm); + if matches!(&result, Ok(PyIterReturn::Return(_))) || !strict.load(atomic::Ordering::Acquire) { + return result; + } + // Native iterators can surface StopIteration as an exception as well as + // the dedicated return variant (for example, a filter item's __bool__). + // Preserve the original exception in non-strict mode. + match result.or_else(|error| PyIterReturn::from_pyresult(Err(error), vm))? { + PyIterReturn::StopIteration(value) => { + drop(value); + strict_exhausted(iterators, index, name, vm) + } + item => Ok(item), + } +} + +// Both map and zip check the remaining sources only when the first one stops. +// Callers clear its StopIteration value before checking: strict exhaustion does +// not propagate source return values, unlike ordinary shortest-input iteration. +fn strict_exhausted( + iterators: &[PyIter], + stopped: usize, + name: &str, + vm: &VirtualMachine, +) -> PyResult { + let (index, direction) = if stopped > 0 { + (stopped, "shorter") + } else { + let mut longer = None; + for (index, iterator) in iterators.iter().enumerate().skip(1) { + if let PyIterReturn::Return(item) = iterator + .next(vm) + .or_else(|error| PyIterReturn::from_pyresult(Err(error), vm))? + { + drop(item); + longer = Some(index); + break; + } + } + let Some(index) = longer else { + return Ok(PyIterReturn::StopIteration(None)); + }; + (index, "longer") + }; + let plural = if index == 1 { " " } else { "s 1-" }; + Err(vm.new_value_error(format!( + "{name}() argument {} is {direction} than argument{plural}{index}", + index + 1, + ))) +} + pub(crate) fn init(ctx: &'static Context) { PyZip::extend_class(ctx, ctx.types.zip_type); } From 2bdbe0427454a64cd509b9ba1210971c0f0322bb Mon Sep 17 00:00:00 2001 From: dot Date: Wed, 4 Nov 2026 20:32:43 +0900 Subject: [PATCH 2/6] Share generic sequence iterator callback controls Keep sequence callbacks outside state locks and match forward/reverse state, exhaustion and signed-wrap behavior without changing native container iterators. Assisted-by: OpenAI Codex:model version not exposed --- crates/vm/src/builtins/enumerate.rs | 37 ++--- crates/vm/src/builtins/iter.rs | 242 ++++++++++++++++++++++++---- 2 files changed, 225 insertions(+), 54 deletions(-) diff --git a/crates/vm/src/builtins/enumerate.rs b/crates/vm/src/builtins/enumerate.rs index 7610d92dbd5..df6270783a9 100644 --- a/crates/vm/src/builtins/enumerate.rs +++ b/crates/vm/src/builtins/enumerate.rs @@ -1,8 +1,8 @@ use super::{ - IterStatus, PositionIterInternal, PyGenericAlias, PyTupleRef, PyType, PyTypeRef, - iter::builtins_reversed, locked_rev_next, + PyGenericAlias, PyTupleRef, PyType, PyTypeRef, + iter::{SequenceIter, SequenceIterDirection, builtins_reversed}, }; -use crate::common::lock::{PyMutex, PyRwLock}; +use crate::common::lock::PyRwLock; use crate::{ AsObject, Context, Py, PyObjectRef, PyPayload, PyResult, VirtualMachine, class::PyClassImpl, @@ -139,7 +139,7 @@ impl IterNext for PyEnumerate { #[pyclass(module = false, name = "reversed", traverse)] #[derive(Debug)] pub(crate) struct PyReverseSequenceIterator { - internal: PyMutex>, + internal: SequenceIter, } impl PyPayload for PyReverseSequenceIterator { @@ -151,9 +151,9 @@ impl PyPayload for PyReverseSequenceIterator { impl PyReverseSequenceIterator { pub(crate) const fn new(obj: PyObjectRef, len: usize) -> Self { - let position = len.saturating_sub(1); + let position = len as isize - 1; Self { - internal: PyMutex::new(PositionIterInternal::new(obj, position)), + internal: SequenceIter::new(obj, position, SequenceIterDirection::Reverse), } } } @@ -162,38 +162,29 @@ impl PyReverseSequenceIterator { impl Py { #[pymethod] fn __length_hint__(&self, vm: &VirtualMachine) -> PyResult { - let internal = self.internal.lock(); - if let IterStatus::Active(obj) = &internal.status - && internal.position <= obj.length(vm)? - { - return Ok(internal.position + 1); - } - Ok(0) + vm.with_recursion(" in __length_hint__", || { + self.internal + .length_hint(vm) + .map(|length| length.unwrap_or(0)) + }) } #[pymethod] fn __setstate__(&self, state: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { - self.internal.lock().set_state(&state, |_, pos| pos, vm) + self.internal.set_state(&state, vm) } #[pymethod] fn __reduce__(&self, vm: &VirtualMachine) -> PyResult { let func = builtins_reversed(vm)?; - Ok(self.internal.lock().reduce( - func, - |x| x.clone(), - |vm| vm.ctx.empty_tuple.clone().into(), - vm, - )) + Ok(self.internal.reduce(func, vm)) } } impl SelfIter for PyReverseSequenceIterator {} impl IterNext for PyReverseSequenceIterator { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { - locked_rev_next(&zelf.internal, |obj, pos| { - PyIterReturn::from_getitem_result(obj.get_item(&pos, vm), vm) - }) + zelf.internal.next(vm) } } diff --git a/crates/vm/src/builtins/iter.rs b/crates/vm/src/builtins/iter.rs index 365aed9da6f..4b99cf3919f 100644 --- a/crates/vm/src/builtins/iter.rs +++ b/crates/vm/src/builtins/iter.rs @@ -224,10 +224,212 @@ pub fn builtins_reversed(vm: &VirtualMachine) -> PyResult { vm.eval_get_builtin(vm.ctx.intern_str("reversed")) } +#[derive(Debug, Clone, Copy)] +pub(crate) enum SequenceIterDirection { + Forward, + Reverse, +} + +#[derive(Debug, Traverse)] +struct SequenceIterState { + sequence: Option, + #[pytraverse(skip)] + index: isize, +} + +/// State for generic iterators whose sequence slots can execute Python code. +/// Unlike the concrete container iterators, these must release their state +/// lock before invoking a slot, allocating a Python object, or dropping an owner. +#[derive(Debug, Traverse)] +pub(crate) struct SequenceIter { + state: PyMutex, + #[pytraverse(skip)] + direction: SequenceIterDirection, +} + +impl SequenceIter { + pub(crate) const fn new( + sequence: PyObjectRef, + index: isize, + direction: SequenceIterDirection, + ) -> Self { + Self { + state: PyMutex::new(SequenceIterState { + sequence: Some(sequence), + index, + }), + direction, + } + } + + fn snapshot(&self) -> (Option, isize) { + let state = self.state.lock(); + (state.sequence.clone(), state.index) + } + + fn exhaust(&self) { + let released = { + let mut state = self.state.lock(); + if matches!(self.direction, SequenceIterDirection::Reverse) { + state.index = -1; + } + state.sequence.take() + }; + drop(released); + } + + fn get_item(sequence: &PyObject, mut index: isize, vm: &VirtualMachine) -> PyResult { + let seq = sequence.sequence_unchecked(); + // A reentrant state change can wrap the forward index below zero. + // PySequence_GetItem adjusts it once using the sequence length slot. + // Keep this separate from the concrete containers' index handling. + if index < 0 + && seq.slots().has_item() + && let Some(length) = seq.length_opt(vm) + { + index = index.wrapping_add(length? as isize); + } + seq.get_item(index, vm) + } + + pub(crate) fn next(&self, vm: &VirtualMachine) -> PyResult { + let (sequence, index) = self.snapshot(); + let Some(sequence) = sequence else { + return Ok(PyIterReturn::StopIteration(None)); + }; + match self.direction { + SequenceIterDirection::Forward if index == isize::MAX => { + drop(sequence); + return Err(vm.new_overflow_error("iter index too large")); + } + SequenceIterDirection::Reverse if index < 0 => { + drop(sequence); + self.exhaust(); + return Ok(PyIterReturn::StopIteration(None)); + } + _ => {} + } + + let result = Self::get_item(&sequence, index, vm); + // A reentrant call can have removed the iterator's sequence owner. + // Release our temporary owner before publishing the outer result. + drop(sequence); + match result { + Ok(item) => { + let mut state = self.state.lock(); + state.index = match self.direction { + // Forward iteration advances the current index, including + // changes made by __getitem__ or a nested next(). + // CPython's default no-strict-overflow build wraps if + // that callback restored PY_SSIZE_T_MAX. + SequenceIterDirection::Forward => state.index.wrapping_add(1), + // Reverse iteration instead advances the captured index. + // Keep the sequence at -1 until the next next() call. + SequenceIterDirection::Reverse => index - 1, + }; + drop(state); + Ok(PyIterReturn::Return(item)) + } + Err(error) + if error.fast_isinstance(vm.ctx.exceptions.index_error) + || error.fast_isinstance(vm.ctx.exceptions.stop_iteration) => + { + // Clear the whole exception, including StopIteration.value, + // before publishing exhaustion. Its finalizer may reenter. + drop(error); + self.exhaust(); + Ok(PyIterReturn::StopIteration(None)) + } + Err(error) => { + if matches!(self.direction, SequenceIterDirection::Reverse) { + self.exhaust(); + } + Err(error) + } + } + } + + pub(crate) fn length_hint(&self, vm: &VirtualMachine) -> PyResult> { + let (sequence, index) = self.snapshot(); + let Some(sequence) = sequence else { + return Ok(Some(0)); + }; + if matches!(self.direction, SequenceIterDirection::Reverse) && index < 0 { + return Ok(Some(0)); + } + let seq = sequence.sequence_unchecked(); + if matches!(self.direction, SequenceIterDirection::Forward) + && seq.slots().length.load().is_none() + && sequence.mapping_unchecked().slots().length.load().is_none() + { + return Ok(None); + } + let length = seq.length(vm)?; + drop(sequence); + let remaining = match self.direction { + SequenceIterDirection::Forward => { + // iter_len reads the index after the length callback. + let index = self.state.lock().index; + (length as isize).wrapping_sub(index).max(0) as usize + } + SequenceIterDirection::Reverse => { + // reversed_len uses the index captured before the callback. + let remaining = index as usize + 1; + if length < remaining { 0 } else { remaining } + } + }; + Ok(Some(remaining)) + } + + pub(crate) fn set_state(&self, state: &PyObject, vm: &VirtualMachine) -> PyResult<()> { + // PyLong_AsSsize_t accepts int subclasses, not __index__. Convert even + // when exhausted, and retain the original argument through callbacks. + let index = state + .downcast_ref::() + .ok_or_else(|| vm.new_type_error("an integer is required"))?; + let index = isize::try_from(index.as_bigint()) + .map_err(|_| vm.new_overflow_error("Python int too large to convert to C ssize_t"))?; + match self.direction { + SequenceIterDirection::Forward => { + let mut state = self.state.lock(); + if state.sequence.is_some() { + state.index = index.max(0); + } + } + SequenceIterDirection::Reverse => { + let (sequence, current) = self.snapshot(); + if current < 0 { + return Ok(()); + } + let Some(sequence) = sequence else { + return Ok(()); + }; + let length = sequence.sequence_unchecked().length(vm)?; + drop(sequence); + self.state.lock().index = index.clamp(-1, length as isize - 1); + } + } + Ok(()) + } + + /// Resolve the reconstruction function before taking this snapshot. Tuple + /// allocation can also run GC and finalizers, so it follows the lock scope. + pub(crate) fn reduce(&self, func: PyObjectRef, vm: &VirtualMachine) -> PyTupleRef { + let (sequence, index) = self.snapshot(); + if let Some(sequence) = sequence + && (matches!(self.direction, SequenceIterDirection::Forward) || index >= 0) + { + vm.new_tuple((func, (sequence,), index)) + } else { + vm.new_tuple((func, (vm.ctx.empty_tuple.clone(),))) + } + } +} + #[pyclass(module = false, name = "iterator", traverse)] #[derive(Debug)] pub struct PySequenceIterator { - internal: PyMutex>, + internal: SequenceIter, } impl PyPayload for PySequenceIterator { @@ -241,7 +443,7 @@ impl PySequenceIterator { pub fn new(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult { let _seq = obj.try_sequence(vm)?; Ok(Self { - internal: PyMutex::new(PositionIterInternal::new(obj, 0)), + internal: SequenceIter::new(obj, 0, SequenceIterDirection::Forward), }) } } @@ -251,51 +453,29 @@ impl Py { #[pymethod] fn __length_hint__(&self, vm: &VirtualMachine) -> PyResult { vm.with_recursion(" in __length_hint__", || { - let (obj, position) = { - let internal = self.internal.lock(); - match &internal.status { - IterStatus::Active(obj) => (Some(obj.clone()), internal.position), - IterStatus::Exhausted => (None, 0), - } - }; - if let Some(obj) = obj { - let seq = obj.sequence_unchecked(); - match seq.length_opt(vm) { - Some(len) => { - len.map(|len| PyInt::from(len.saturating_sub(position)).into_pyobject(vm)) - } - None => Ok(vm.ctx.not_implemented()), - } - } else { - Ok(PyInt::from(0).into_pyobject(vm)) - } + Ok(match self.internal.length_hint(vm)? { + Some(length) => PyInt::from(length).into_pyobject(vm), + None => vm.ctx.not_implemented(), + }) }) } #[pymethod] fn __reduce__(&self, vm: &VirtualMachine) -> PyResult { let func = builtins_iter(vm)?; - Ok(self.internal.lock().reduce( - func, - |x| x.clone(), - |vm| vm.ctx.empty_tuple.clone().into(), - vm, - )) + Ok(self.internal.reduce(func, vm)) } #[pymethod] fn __setstate__(&self, state: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { - self.internal.lock().set_state(&state, |_, pos| pos, vm) + self.internal.set_state(&state, vm) } } impl SelfIter for PySequenceIterator {} impl IterNext for PySequenceIterator { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { - locked_next(&zelf.internal, |obj, pos| { - let seq = obj.sequence_unchecked(); - PyIterReturn::from_getitem_result(seq.get_item(pos as isize, vm), vm) - }) + zelf.internal.next(vm) } } From 3ac5895566a8767cbd487d1141ca2ba843bab695 Mon Sep 17 00:00:00 2001 From: dot Date: Wed, 4 Nov 2026 21:24:30 +0900 Subject: [PATCH 3/6] Share strict iterator ssize conversion Assisted-by: OpenAI Codex:model version not exposed --- crates/vm/src/builtins/iter.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/crates/vm/src/builtins/iter.rs b/crates/vm/src/builtins/iter.rs index 4b99cf3919f..6ddec7b311d 100644 --- a/crates/vm/src/builtins/iter.rs +++ b/crates/vm/src/builtins/iter.rs @@ -6,7 +6,7 @@ use super::{PyInt, PyTupleRef, PyType}; use crate::{ AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyResult, VirtualMachine, class::PyClassImpl, - function::ArgCallable, + function::{ArgCallable, int_as_ssize}, object::{Traverse, TraverseFn}, protocol::PyIterReturn, types::{IterNext, Iterable, PyComparisonOp, SelfIter}, @@ -384,11 +384,7 @@ impl SequenceIter { pub(crate) fn set_state(&self, state: &PyObject, vm: &VirtualMachine) -> PyResult<()> { // PyLong_AsSsize_t accepts int subclasses, not __index__. Convert even // when exhausted, and retain the original argument through callbacks. - let index = state - .downcast_ref::() - .ok_or_else(|| vm.new_type_error("an integer is required"))?; - let index = isize::try_from(index.as_bigint()) - .map_err(|_| vm.new_overflow_error("Python int too large to convert to C ssize_t"))?; + let index = int_as_ssize(state, vm)?; match self.direction { SequenceIterDirection::Forward => { let mut state = self.state.lock(); From 163c697d12012af38e920df5aab2e99cd46ca02c Mon Sep 17 00:00:00 2001 From: dot Date: Wed, 4 Nov 2026 22:37:54 +0900 Subject: [PATCH 4/6] Align reversed construction and proxy dispatch Assisted-by: OpenAI Codex:model version not exposed --- crates/vm/src/builtins/enumerate.rs | 90 +++++++++++++++++++++++++++-- crates/vm/src/builtins/weakproxy.rs | 3 +- crates/vm/src/stdlib/builtins.rs | 21 +------ 3 files changed, 88 insertions(+), 26 deletions(-) diff --git a/crates/vm/src/builtins/enumerate.rs b/crates/vm/src/builtins/enumerate.rs index df6270783a9..99a6ab9f559 100644 --- a/crates/vm/src/builtins/enumerate.rs +++ b/crates/vm/src/builtins/enumerate.rs @@ -1,15 +1,15 @@ use super::{ PyGenericAlias, PyTupleRef, PyType, PyTypeRef, - iter::{SequenceIter, SequenceIterDirection, builtins_reversed}, + iter::{SequenceIter, SequenceIterDirection}, }; use crate::common::lock::PyRwLock; use crate::{ AsObject, Context, Py, PyObjectRef, PyPayload, PyResult, VirtualMachine, class::PyClassImpl, - function::ArgIndex, + function::{ArgIndex, FuncArgs}, protocol::{PyIter, PyIterReturn}, raise_if_stop, - types::{Constructor, IterNext, Iterable, SelfIter}, + types::{Constructor, IterNext, Iterable, SelfIter, drop_kwargs_if_init_overridden}, }; use malachite_bigint::BigInt; use num_traits::ToPrimitive; @@ -150,6 +150,44 @@ impl PyPayload for PyReverseSequenceIterator { } impl PyReverseSequenceIterator { + pub(crate) fn from_object( + cls: PyTypeRef, + sequence: &PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + // Keep the original sequence alive while binding, calling, and releasing + // the special method, including when it does not retain its receiver. + if let Some(method) = vm.get_method(sequence.clone(), identifier!(vm, __reversed__)) { + let method = method?; + if vm.is_none(&method) { + return Err(Self::not_reversible(sequence, vm)); + } + let result = method.call((), vm); + drop(method); + return result; + } + + Self::py_new( + &cls, + ReversedArgs { + object: sequence.clone(), + }, + vm, + )? + .into_ref_with_type(vm, cls) + .map(Into::into) + } + + fn not_reversible( + sequence: &PyObjectRef, + vm: &VirtualMachine, + ) -> crate::builtins::PyBaseExceptionRef { + vm.new_type_error(format!( + "'{}' object is not reversible", + sequence.class().slot_name() + )) + } + pub(crate) const fn new(obj: PyObjectRef, len: usize) -> Self { let position = len as isize - 1; Self { @@ -158,7 +196,49 @@ impl PyReverseSequenceIterator { } } -#[pyclass(with(IterNext, Iterable))] +#[derive(FromArgs)] +pub(crate) struct ReversedArgs { + #[pyarg(positional)] + object: PyObjectRef, +} + +impl Constructor for PyReverseSequenceIterator { + type Args = ReversedArgs; + + fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult { + // A subtype's custom initializer may consume keywords, but the exact + // type and subtypes inheriting its initializer reject them first. + // Retain the original arguments through callbacks even if keywords + // are removed from the constructor's view of the call. + let constructor_args = + drop_kwargs_if_init_overridden(&cls, vm.ctx.types.reverse_iter_type, args.clone()); + if !constructor_args.kwargs.is_empty() { + return Err(vm.new_type_error("reversed() takes no keyword arguments")); + } + let ReversedArgs { object } = constructor_args.bind_for(vm, "reversed")?; + let result = Self::from_object(cls, &object, vm); + drop(object); + drop(args); + result + } + + fn py_new( + _cls: &Py, + Self::Args { object }: Self::Args, + vm: &VirtualMachine, + ) -> PyResult { + // The slot handles arbitrary __reversed__ results before constructing + // a payload. Checking sequence support must not bind __getitem__. + let protocol = object.sequence_unchecked(); + if !protocol.check() { + return Err(Self::not_reversible(&object, vm)); + } + let len = protocol.length(vm)?; + Ok(Self::new(object, len)) + } +} + +#[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE))] impl Py { #[pymethod] fn __length_hint__(&self, vm: &VirtualMachine) -> PyResult { @@ -176,7 +256,7 @@ impl Py { #[pymethod] fn __reduce__(&self, vm: &VirtualMachine) -> PyResult { - let func = builtins_reversed(vm)?; + let func = self.class().to_owned().into(); Ok(self.internal.reduce(func, vm)) } } diff --git a/crates/vm/src/builtins/weakproxy.rs b/crates/vm/src/builtins/weakproxy.rs index d0afee4f254..790d9ab847e 100644 --- a/crates/vm/src/builtins/weakproxy.rs +++ b/crates/vm/src/builtins/weakproxy.rs @@ -7,7 +7,6 @@ use crate::{ common::hash::PyHash, function::{Either, FuncArgs, PyArithmeticValue, PyComparisonValue, PySetterValue}, protocol::{PyIter, PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods}, - stdlib::builtins::reversed, types::{ AsMapping, AsNumber, AsSequence, Callable, Comparable, Constructor, GetAttr, Hashable, IterNext, Iterable, PyComparisonOp, Representable, SetAttr, @@ -190,7 +189,7 @@ impl Py { #[pymethod] fn __reversed__(&self, vm: &VirtualMachine) -> PyResult { let obj = self.try_upgrade(vm)?; - reversed(obj, vm) + vm.call_method(&obj, identifier!(vm, __reversed__).as_str(), ()) } } diff --git a/crates/vm/src/stdlib/builtins.rs b/crates/vm/src/stdlib/builtins.rs index e377fb1d3cb..642e673fc78 100644 --- a/crates/vm/src/stdlib/builtins.rs +++ b/crates/vm/src/stdlib/builtins.rs @@ -3,7 +3,7 @@ //! Implements the list of [builtin Python functions](https://docs.python.org/3/library/builtins.html). use crate::{Py, VirtualMachine, builtins::PyModule, class::PyClassImpl}; pub(crate) use builtins::{DOC, module_def}; -pub use builtins::{ascii, print, reversed}; +pub use builtins::{ascii, print}; #[pymodule] mod builtins { @@ -13,7 +13,6 @@ mod builtins { builtins::{ PyAnyDict, PyByteArray, PyBytes, PyComplex, PyFloat, PyInt, PyStr, PyStrRef, PyTuple, PyTupleRef, PyType, PyUtf8StrRef, - enumerate::PyReverseSequenceIterator, function::{PyCell, PyCellRef, PyFunction}, iter::PyCallableIterator, list::{PyList, SortOptions}, @@ -1199,23 +1198,6 @@ mod builtins { obj.repr(vm) } - #[pyfunction] - pub fn reversed(object: PyObjectRef, vm: &VirtualMachine) -> PyResult { - let sequence = object; - if let Some(reversed_method) = - vm.get_method(sequence.clone(), identifier!(vm, __reversed__)) - { - reversed_method?.call((), vm) - } else { - vm.get_method_or_type_error(sequence.clone(), identifier!(vm, __getitem__), || { - "argument to reversed() must be a sequence".to_owned() - })?; - let len = sequence.length(vm)?; - let obj_iterator = PyReverseSequenceIterator::new(sequence, len); - Ok(obj_iterator.into_pyobject(vm)) - } - } - #[derive(FromArgs)] pub(super) struct RoundArgs { number: PyObjectRef, @@ -1695,6 +1677,7 @@ pub fn init_module(vm: &VirtualMachine, module: &Py) { "complex" => ctx.types.complex_type.to_owned(), "dict" => ctx.types.dict_type.to_owned(), "enumerate" => ctx.types.enumerate_type.to_owned(), + "reversed" => ctx.types.reverse_iter_type.to_owned(), "float" => ctx.types.float_type.to_owned(), "frozendict" => ctx.types.frozendict_type.to_owned(), "frozenset" => ctx.types.frozenset_type.to_owned(), From 26738823b0ce45914d9bc25cac5a9807623888f0 Mon Sep 17 00:00:00 2001 From: dot Date: Wed, 4 Nov 2026 22:43:13 +0900 Subject: [PATCH 5/6] Return reversed reduction without redundant result wrapper Assisted-by: OpenAI Codex:model version not exposed --- crates/vm/src/builtins/enumerate.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/vm/src/builtins/enumerate.rs b/crates/vm/src/builtins/enumerate.rs index 99a6ab9f559..0231c1ed184 100644 --- a/crates/vm/src/builtins/enumerate.rs +++ b/crates/vm/src/builtins/enumerate.rs @@ -255,9 +255,9 @@ impl Py { } #[pymethod] - fn __reduce__(&self, vm: &VirtualMachine) -> PyResult { + fn __reduce__(&self, vm: &VirtualMachine) -> PyTupleRef { let func = self.class().to_owned().into(); - Ok(self.internal.reduce(func, vm)) + self.internal.reduce(func, vm) } } From 46ff99a15d5bbb75f9a6df4daf3dc7052394c442 Mon Sep 17 00:00:00 2001 From: dot Date: Thu, 5 Nov 2026 11:56:09 +0900 Subject: [PATCH 6/6] Return integers from byte sequence item slots Generic reversed iteration uses the sequence item slot. Match byte indexing and CPython by returning integer values for bytes and bytearray rather than single-byte bytes objects. Assisted-by: OpenAI Codex:model version not exposed --- crates/vm/src/builtins/bytearray.rs | 2 +- crates/vm/src/builtins/bytes.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/vm/src/builtins/bytearray.rs b/crates/vm/src/builtins/bytearray.rs index bbcca3a7b86..c90c140adfb 100644 --- a/crates/vm/src/builtins/bytearray.rs +++ b/crates/vm/src/builtins/bytearray.rs @@ -1504,7 +1504,7 @@ impl AsSequence for PyByteArray { PyByteArray::sequence_downcast(seq) .borrow_buf() .getitem_by_index(vm, i) - .map(|x| vm.ctx.new_bytes(vec![x]).into()) + .map(|x| vm.ctx.new_int(x).into()) }), ass_item: atomic_func!(|seq, i, value, vm| { let zelf = PyByteArray::sequence_downcast(seq); diff --git a/crates/vm/src/builtins/bytes.rs b/crates/vm/src/builtins/bytes.rs index d8e95b2a371..239b146536e 100644 --- a/crates/vm/src/builtins/bytes.rs +++ b/crates/vm/src/builtins/bytes.rs @@ -698,7 +698,7 @@ impl AsSequence for PyBytes { PyBytes::sequence_downcast(seq) .as_bytes() .getitem_by_index(vm, i) - .map(|x| vm.ctx.new_bytes(vec![x]).into()) + .map(|x| vm.ctx.new_int(x).into()) }), contains: atomic_func!(|seq, other, vm| { PyBytes::sequence_downcast(seq).__contains__(other.to_owned(), vm)