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) diff --git a/crates/vm/src/builtins/enumerate.rs b/crates/vm/src/builtins/enumerate.rs index 7610d92dbd5..0231c1ed184 100644 --- a/crates/vm/src/builtins/enumerate.rs +++ b/crates/vm/src/builtins/enumerate.rs @@ -1,15 +1,15 @@ use super::{ - IterStatus, PositionIterInternal, PyGenericAlias, PyTupleRef, PyType, PyTypeRef, - iter::builtins_reversed, locked_rev_next, + PyGenericAlias, PyTupleRef, PyType, PyTypeRef, + iter::{SequenceIter, SequenceIterDirection}, }; -use crate::common::lock::{PyMutex, PyRwLock}; +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; @@ -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 { @@ -150,50 +150,121 @@ 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.saturating_sub(1); + let position = len as isize - 1; Self { - internal: PyMutex::new(PositionIterInternal::new(obj, position)), + internal: SequenceIter::new(obj, position, SequenceIterDirection::Reverse), } } } -#[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 { - 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, - )) + fn __reduce__(&self, vm: &VirtualMachine) -> PyTupleRef { + let func = self.class().to_owned().into(); + 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/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..6ddec7b311d 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, + function::{ArgCallable, int_as_ssize}, 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)] @@ -224,10 +224,208 @@ 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 = int_as_ssize(state, vm)?; + 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 +439,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,59 +449,36 @@ 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) } } #[pyclass(module = false, name = "callable_iterator", traverse)] #[derive(Debug)] pub struct PyCallableIterator { - sentinel: PyObjectRef, - status: PyRwLock>, + status: PyRwLock>, } impl PyPayload for PyCallableIterator { @@ -317,10 +492,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 +512,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 +539,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/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/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); } 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(),