diff --git a/extra_tests/snippets/builtin_filter.py b/extra_tests/snippets/builtin_filter.py index d0b5ccd5c..25c00803c 100644 --- a/extra_tests/snippets/builtin_filter.py +++ b/extra_tests/snippets/builtin_filter.py @@ -29,4 +29,5 @@ def predicate(x): return True -assert list(filter(predicate, [1, 2, 0, 4, 5])) == [1, 2] +filtered = list(filter(predicate, [1, 2, 0, 4, 5])) +assert filtered == [1, 2], filtered diff --git a/stdlib/src/array.rs b/stdlib/src/array.rs index 262ea06f5..29a675e11 100644 --- a/stdlib/src/array.rs +++ b/stdlib/src/array.rs @@ -17,7 +17,7 @@ mod array { }, class_or_notimplemented, function::{ArgBytesLike, ArgIterable, OptionalArg}, - protocol::{BufferInternal, BufferOptions, PyBuffer, ResizeGuard}, + protocol::{BufferInternal, BufferOptions, PyBuffer, PyIterReturn, ResizeGuard}, sliceable::{saturate_index, PySliceableSequence, PySliceableSequenceMut, SequenceIndex}, slots::{ AsBuffer, Comparable, Iterable, IteratorIterable, PyComparisonOp, SlotConstructor, @@ -1199,13 +1199,14 @@ mod array { impl IteratorIterable for PyArrayIter {} impl SlotIterator for PyArrayIter { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let pos = zelf.position.fetch_add(1); - if let Some(item) = zelf.array.read().getitem_by_idx(pos, vm)? { - Ok(item) + let r = if let Some(item) = zelf.array.read().getitem_by_idx(pos, vm)? { + PyIterReturn::Return(item) } else { - Err(vm.new_stop_iteration()) - } + PyIterReturn::StopIteration(None) + }; + Ok(r) } } diff --git a/stdlib/src/csv.rs b/stdlib/src/csv.rs index 462df4c79..03919ee5b 100644 --- a/stdlib/src/csv.rs +++ b/stdlib/src/csv.rs @@ -3,7 +3,7 @@ use crate::vm::{ builtins::{PyStr, PyStrRef}, function::{ArgIterable, ArgumentError, FromArgs, FuncArgs}, match_class, named_function, - protocol::PyIter, + protocol::{PyIter, PyIterReturn}, py_module, slots::{IteratorIterable, SlotIterator}, types::create_simple_type, @@ -101,8 +101,11 @@ impl fmt::Debug for Reader { impl Reader {} impl IteratorIterable for Reader {} impl SlotIterator for Reader { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - let string = zelf.iter.next(vm)?; + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let string = match zelf.iter.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + }; let string = string.downcast::().map_err(|obj| { vm.new_type_error(format!( "iterator should return strings, not {} (the file should be opened in text mode)", @@ -136,7 +139,7 @@ impl SlotIterator for Reader { csv_core::ReadRecordResult::OutputFull => resize_buf(buffer), csv_core::ReadRecordResult::OutputEndsFull => resize_buf(output_ends), csv_core::ReadRecordResult::Record => break, - csv_core::ReadRecordResult::End => return Err(vm.new_stop_iteration()), + csv_core::ReadRecordResult::End => return Ok(PyIterReturn::StopIteration(None)), } } let rest = &input[input_offset..]; @@ -160,7 +163,7 @@ impl SlotIterator for Reader { .map_err(|_e| vm.new_unicode_decode_error("csv not utf8".to_owned())) }) .collect::>()?; - Ok(vm.ctx.new_list(out)) + Ok(PyIterReturn::Return(vm.ctx.new_list(out))) } } diff --git a/vm/src/builtins/asyncgenerator.rs b/vm/src/builtins/asyncgenerator.rs index 26a8eb7f7..ae6a3ad41 100644 --- a/vm/src/builtins/asyncgenerator.rs +++ b/vm/src/builtins/asyncgenerator.rs @@ -4,6 +4,7 @@ use crate::{ coroutine::{Coro, Variant}, frame::FrameRef, function::OptionalArg, + protocol::PyIterReturn, slots::{IteratorIterable, SlotIterator}, IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol, VirtualMachine, @@ -256,8 +257,9 @@ impl PyAsyncGenASend { impl IteratorIterable for PyAsyncGenASend {} impl SlotIterator for PyAsyncGenASend { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - zelf.send(vm.ctx.none(), vm) + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + // TODO: Fix zelf.send to return PyIterReturn + PyIterReturn::from_result(zelf.send(vm.ctx.none(), vm), vm) } } @@ -398,8 +400,9 @@ impl PyAsyncGenAThrow { impl IteratorIterable for PyAsyncGenAThrow {} impl SlotIterator for PyAsyncGenAThrow { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - zelf.send(vm.ctx.none(), vm) + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + // TODO: Fix zelf.send to return PyIterReturn + PyIterReturn::from_result(zelf.send(vm.ctx.none(), vm), vm) } } diff --git a/vm/src/builtins/bytearray.rs b/vm/src/builtins/bytearray.rs index 7abd30bfb..2ab6091dc 100644 --- a/vm/src/builtins/bytearray.rs +++ b/vm/src/builtins/bytearray.rs @@ -15,7 +15,7 @@ use crate::{ ByteInnerTranslateOptions, DecodeArgs, PyBytesInner, }, function::{ArgBytesLike, ArgIterable, FuncArgs, OptionalArg, OptionalOption}, - protocol::{BufferInternal, BufferOptions, PyBuffer, ResizeGuard}, + protocol::{BufferInternal, BufferOptions, PyBuffer, PyIterReturn, ResizeGuard}, sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex}, slots::{ AsBuffer, Callable, Comparable, Hashable, Iterable, IteratorIterable, PyComparisonOp, @@ -745,12 +745,13 @@ impl PyValue for PyByteArrayIterator { impl PyByteArrayIterator {} impl IteratorIterable for PyByteArrayIterator {} impl SlotIterator for PyByteArrayIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let pos = zelf.position.fetch_add(1); - if let Some(&ret) = zelf.bytearray.borrow_buf().get(pos) { - Ok(ret.into_pyobject(vm)) + let r = if let Some(&ret) = zelf.bytearray.borrow_buf().get(pos) { + PyIterReturn::Return(ret.into_pyobject(vm)) } else { - Err(vm.new_stop_iteration()) - } + PyIterReturn::StopIteration(None) + }; + Ok(r) } } diff --git a/vm/src/builtins/bytes.rs b/vm/src/builtins/bytes.rs index ed6ea9089..2df600017 100644 --- a/vm/src/builtins/bytes.rs +++ b/vm/src/builtins/bytes.rs @@ -7,7 +7,7 @@ use crate::{ }, common::hash::PyHash, function::{ArgBytesLike, ArgIterable, OptionalArg, OptionalOption}, - protocol::{BufferInternal, BufferOptions, PyBuffer}, + protocol::{BufferInternal, BufferOptions, PyBuffer, PyIterReturn}, slots::{ AsBuffer, Callable, Comparable, Hashable, Iterable, IteratorIterable, PyComparisonOp, SlotConstructor, SlotIterator, @@ -598,13 +598,14 @@ impl PyValue for PyBytesIterator { impl PyBytesIterator {} impl IteratorIterable for PyBytesIterator {} impl SlotIterator for PyBytesIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let pos = zelf.position.fetch_add(1); - if let Some(&ret) = zelf.bytes.as_bytes().get(pos) { - Ok(vm.ctx.new_int(ret)) + let r = if let Some(&ret) = zelf.bytes.as_bytes().get(pos) { + PyIterReturn::Return(vm.ctx.new_int(ret)) } else { - Err(vm.new_stop_iteration()) - } + PyIterReturn::StopIteration(None) + }; + Ok(r) } } diff --git a/vm/src/builtins/coroutine.rs b/vm/src/builtins/coroutine.rs index c3b36a86c..ccf4016a4 100644 --- a/vm/src/builtins/coroutine.rs +++ b/vm/src/builtins/coroutine.rs @@ -3,6 +3,7 @@ use crate::{ coroutine::{Coro, Variant}, frame::FrameRef, function::OptionalArg, + protocol::PyIterReturn, slots::{IteratorIterable, SlotIterator}, IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, VirtualMachine, }; @@ -103,8 +104,9 @@ impl PyCoroutine { impl IteratorIterable for PyCoroutine {} impl SlotIterator for PyCoroutine { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - zelf.send(vm.ctx.none(), vm) + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + // TODO: Fix zelf.send to return PyIterReturn + PyIterReturn::from_result(zelf.send(vm.ctx.none(), vm), vm) } } @@ -141,8 +143,9 @@ impl PyCoroutineWrapper { impl IteratorIterable for PyCoroutineWrapper {} impl SlotIterator for PyCoroutineWrapper { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - zelf.send(vm.ctx.none(), vm) + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + // TODO: Fix zelf.send to return PyIterReturn + PyIterReturn::from_result(zelf.send(vm.ctx.none(), vm), vm) } } diff --git a/vm/src/builtins/dict.rs b/vm/src/builtins/dict.rs index e9b5b549b..7a57d688d 100644 --- a/vm/src/builtins/dict.rs +++ b/vm/src/builtins/dict.rs @@ -4,7 +4,7 @@ use crate::{ common::ascii, dictdatatype::{self, DictKey}, function::{ArgIterable, FuncArgs, KwArgs, OptionalArg}, - iterator, + protocol::PyIterReturn, slots::{ Comparable, Hashable, Iterable, IteratorIterable, PyComparisonOp, SlotIterator, Unhashable, }, @@ -86,7 +86,7 @@ impl PyDict { } } else if let Some(keys) = vm.get_method(dict_obj.clone(), "keys") { let keys = vm.invoke(&keys?, ())?.get_iter(vm)?; - while let Some(key) = iterator::get_next_object(vm, &keys)? { + while let PyIterReturn::Return(key) = keys.next(vm)? { let val = dict_obj.get_item(key.clone(), vm)?; dict.insert(vm, key, val)?; } @@ -96,15 +96,20 @@ impl PyDict { fn err(vm: &VirtualMachine) -> PyBaseExceptionRef { vm.new_value_error("Iterator must have exactly two elements".to_owned()) } - let element = match iterator::get_next_object(vm, &iter)? { - Some(obj) => obj, - None => break, + let element = match iter.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(_) => break, }; let elem_iter = element.get_iter(vm)?; - let key = iterator::get_next_object(vm, &elem_iter)?.ok_or_else(|| err(vm))?; - let value = - iterator::get_next_object(vm, &elem_iter)?.ok_or_else(|| err(vm))?; - if iterator::get_next_object(vm, &elem_iter)?.is_some() { + let key = match elem_iter.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(_) => return Err(err(vm)), + }; + let value = match elem_iter.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(_) => return Err(err(vm)), + }; + if matches!(elem_iter.next(vm)?, PyIterReturn::Return(_)) { return Err(err(vm)); } dict.insert(vm, key, value)?; @@ -734,9 +739,9 @@ macro_rules! dict_iterator { impl IteratorIterable for $iter_name {} impl SlotIterator for $iter_name { #[allow(clippy::redundant_closure_call)] - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { match zelf.status.load() { - IterStatus::Exhausted => Err(vm.new_stop_iteration()), + IterStatus::Exhausted => Ok(PyIterReturn::StopIteration(None)), IterStatus::Active => { if zelf.dict.entries.has_changed_size(&zelf.size) { zelf.status.store(IterStatus::Exhausted); @@ -745,10 +750,12 @@ macro_rules! dict_iterator { )); } match zelf.dict.entries.next_entry_atomic(&zelf.position) { - Some((key, value)) => Ok(($result_fn)(vm, key, value)), + Some((key, value)) => { + Ok(PyIterReturn::Return(($result_fn)(vm, key, value))) + } None => { zelf.status.store(IterStatus::Exhausted); - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } @@ -795,9 +802,9 @@ macro_rules! dict_iterator { impl IteratorIterable for $reverse_iter_name {} impl SlotIterator for $reverse_iter_name { #[allow(clippy::redundant_closure_call)] - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { match zelf.status.load() { - IterStatus::Exhausted => Err(vm.new_stop_iteration()), + IterStatus::Exhausted => Ok(PyIterReturn::StopIteration(None)), IterStatus::Active => { if zelf.dict.entries.has_changed_size(&zelf.size) { zelf.status.store(IterStatus::Exhausted); @@ -806,10 +813,12 @@ macro_rules! dict_iterator { )); } match zelf.dict.entries.next_entry_atomic_reversed(&zelf.position) { - Some((key, value)) => Ok(($result_fn)(vm, key, value)), + Some((key, value)) => { + Ok(PyIterReturn::Return(($result_fn)(vm, key, value))) + } None => { zelf.status.store(IterStatus::Exhausted); - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } diff --git a/vm/src/builtins/enumerate.rs b/vm/src/builtins/enumerate.rs index 76bf4bcce..e2cfc1bf4 100644 --- a/vm/src/builtins/enumerate.rs +++ b/vm/src/builtins/enumerate.rs @@ -6,7 +6,7 @@ use super::{ use crate::common::lock::PyRwLock; use crate::{ function::OptionalArg, - protocol::PyIter, + protocol::{PyIter, PyIterReturn}, slots::{IteratorIterable, SlotConstructor, SlotIterator}, IntoPyObject, ItemProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol, VirtualMachine, @@ -57,12 +57,15 @@ impl PyEnumerate {} impl IteratorIterable for PyEnumerate {} impl SlotIterator for PyEnumerate { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - let next_obj = zelf.iterator.next(vm)?; + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let next_obj = match zelf.iterator.next(vm)? { + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + PyIterReturn::Return(obj) => obj, + }; let mut counter = zelf.counter.write(); let position = counter.clone(); *counter += 1; - Ok((position, next_obj).into_pyobject(vm)) + Ok(PyIterReturn::Return((position, next_obj).into_pyobject(vm))) } } @@ -139,21 +142,23 @@ impl PyReverseSequenceIterator { impl IteratorIterable for PyReverseSequenceIterator {} impl SlotIterator for PyReverseSequenceIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { if let Exhausted = zelf.status.load() { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } let pos = zelf.position.fetch_sub(1); if pos == 0 { zelf.status.store(Exhausted); } match zelf.obj.get_item(pos, vm) { - Err(ref e) if e.isinstance(&vm.ctx.exceptions.index_error) => { + Err(ref e) + if e.isinstance(&vm.ctx.exceptions.index_error) + || e.isinstance(&vm.ctx.exceptions.stop_iteration) => + { zelf.status.store(Exhausted); - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } - // also catches stop_iteration => stop_iteration - ret => ret, + other => other.map(PyIterReturn::Return), } } } diff --git a/vm/src/builtins/filter.rs b/vm/src/builtins/filter.rs index 802605b38..88e8bf9f2 100644 --- a/vm/src/builtins/filter.rs +++ b/vm/src/builtins/filter.rs @@ -1,6 +1,6 @@ use super::PyTypeRef; use crate::{ - protocol::PyIter, + protocol::{PyIter, PyIterReturn}, slots::{IteratorIterable, SlotConstructor, SlotIterator}, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, VirtualMachine, }; @@ -39,19 +39,25 @@ impl PyFilter {} impl IteratorIterable for PyFilter {} impl SlotIterator for PyFilter { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let predicate = &zelf.predicate; loop { - let next_obj = zelf.iterator.next(vm)?; + let next_obj = match zelf.iterator.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + }; let predicate_value = if vm.is_none(predicate) { next_obj.clone() } else { // the predicate itself can raise StopIteration which does stop the filter // iteration - vm.invoke(predicate, vec![next_obj.clone()])? + match PyIterReturn::from_result(vm.invoke(predicate, vec![next_obj.clone()]), vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + } }; if predicate_value.try_to_bool(vm)? { - return Ok(next_obj); + return Ok(PyIterReturn::Return(next_obj)); } } } diff --git a/vm/src/builtins/generator.rs b/vm/src/builtins/generator.rs index cd1de09d9..deb17e518 100644 --- a/vm/src/builtins/generator.rs +++ b/vm/src/builtins/generator.rs @@ -7,6 +7,7 @@ use crate::{ coroutine::{Coro, Variant}, frame::FrameRef, function::OptionalArg, + protocol::PyIterReturn, slots::{IteratorIterable, SlotIterator}, IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, VirtualMachine, }; @@ -96,8 +97,9 @@ impl PyGenerator { impl IteratorIterable for PyGenerator {} impl SlotIterator for PyGenerator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - zelf.send(vm.ctx.none(), vm) + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + // TODO: Fix zelf.send to return PyIterReturn + PyIterReturn::from_result(zelf.send(vm.ctx.none(), vm), vm) } } diff --git a/vm/src/builtins/iter.rs b/vm/src/builtins/iter.rs index a3e386fa3..efb0756b2 100644 --- a/vm/src/builtins/iter.rs +++ b/vm/src/builtins/iter.rs @@ -5,6 +5,7 @@ use super::{int, PyInt, PyTypeRef}; use crate::{ function::ArgCallable, + protocol::PyIterReturn, slots::{IteratorIterable, SlotIterator}, ItemProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol, VirtualMachine, @@ -92,18 +93,20 @@ impl PySequenceIterator { impl IteratorIterable for PySequenceIterator {} impl SlotIterator for PySequenceIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { if let IterStatus::Exhausted = zelf.status.load() { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } let pos = zelf.position.fetch_add(1); match zelf.obj.get_item(pos, vm) { - Err(ref e) if e.isinstance(&vm.ctx.exceptions.index_error) => { + Err(ref e) + if e.isinstance(&vm.ctx.exceptions.index_error) + || e.isinstance(&vm.ctx.exceptions.stop_iteration) => + { zelf.status.store(IterStatus::Exhausted); - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } - // also catches stop_iteration => stop_iteration - ret => ret, + ret => ret.map(PyIterReturn::Return), } } } @@ -135,16 +138,16 @@ impl PyCallableIterator { impl IteratorIterable for PyCallableIterator {} impl SlotIterator for PyCallableIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { if let IterStatus::Exhausted = zelf.status.load() { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } let ret = zelf.callable.invoke((), vm)?; if vm.bool_eq(&ret, &zelf.sentinel)? { zelf.status.store(IterStatus::Exhausted); - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } else { - Ok(ret) + Ok(PyIterReturn::Return(ret)) } } } diff --git a/vm/src/builtins/list.rs b/vm/src/builtins/list.rs index 11735a5fd..1f0fdfad5 100644 --- a/vm/src/builtins/list.rs +++ b/vm/src/builtins/list.rs @@ -8,6 +8,7 @@ use crate::common::lock::{ }; use crate::{ function::{ArgIterable, FuncArgs, OptionalArg}, + protocol::PyIterReturn, sequence::{self, SimpleSeq}, sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex}, slots::{ @@ -524,17 +525,17 @@ impl PyListIterator { impl IteratorIterable for PyListIterator {} impl SlotIterator for PyListIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, _vm: &VirtualMachine) -> PyResult { if let Exhausted = zelf.status.load() { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } let list = zelf.list.borrow_vec(); let pos = zelf.position.fetch_add(1); if let Some(obj) = list.get(pos) { - Ok(obj.clone()) + Ok(PyIterReturn::Return(obj.clone())) } else { zelf.status.store(Exhausted); - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } @@ -598,15 +599,15 @@ impl PyListReverseIterator { impl IteratorIterable for PyListReverseIterator {} impl SlotIterator for PyListReverseIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, _vm: &VirtualMachine) -> PyResult { if let Exhausted = zelf.status.load() { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } let list = zelf.list.borrow_vec(); let pos = zelf.position.fetch_sub(1); if pos > 0 { if let Some(obj) = list.get(pos) { - return Ok(obj.clone()); + return Ok(PyIterReturn::Return(obj.clone())); } } // We either are == 0 or list.get returned None. Either way, set status @@ -614,10 +615,10 @@ impl SlotIterator for PyListReverseIterator { zelf.status.store(Exhausted); if pos == 0 { if let Some(obj) = list.get(pos) { - return Ok(obj.clone()); + return Ok(PyIterReturn::Return(obj.clone())); } } - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } diff --git a/vm/src/builtins/map.rs b/vm/src/builtins/map.rs index 7d7b174f8..870cb90d3 100644 --- a/vm/src/builtins/map.rs +++ b/vm/src/builtins/map.rs @@ -2,7 +2,7 @@ use super::PyTypeRef; use crate::{ function::PosArgs, iterator, - protocol::PyIter, + protocol::{PyIter, PyIterReturn}, slots::{IteratorIterable, SlotConstructor, SlotIterator}, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, VirtualMachine, }; @@ -47,15 +47,18 @@ impl PyMap { impl IteratorIterable for PyMap {} impl SlotIterator for PyMap { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - let next_objs = zelf - .iterators - .iter() - .map(|iterator| iterator.next(vm)) - .collect::, _>>()?; + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let mut next_objs = Vec::new(); + for iterator in zelf.iterators.iter() { + let item = match iterator.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + }; + next_objs.push(item); + } // the mapper itself can raise StopIteration which does stop the map iteration - vm.invoke(&zelf.mapper, next_objs) + PyIterReturn::from_result(vm.invoke(&zelf.mapper, next_objs), vm) } } diff --git a/vm/src/builtins/pystr.rs b/vm/src/builtins/pystr.rs index 38be33ecd..a3ff1d0bd 100644 --- a/vm/src/builtins/pystr.rs +++ b/vm/src/builtins/pystr.rs @@ -8,6 +8,7 @@ use crate::{ exceptions::IntoPyException, format::{FormatSpec, FormatString, FromTemplate}, function::{ArgIterable, FuncArgs, OptionalArg, OptionalOption}, + protocol::PyIterReturn, sliceable::PySliceableSequence, slots::{ Comparable, Hashable, Iterable, IteratorIterable, PyComparisonOp, SlotConstructor, @@ -230,21 +231,24 @@ impl PyStrIterator { impl IteratorIterable for PyStrIterator {} impl SlotIterator for PyStrIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { if let Exhausted = zelf.status.load() { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } let value = &*zelf.string.as_str(); let mut start = zelf.position.load(atomic::Ordering::SeqCst); loop { if start == value.len() { zelf.status.store(Exhausted); - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } - let ch = value[start..].chars().next().ok_or_else(|| { - zelf.status.store(Exhausted); - vm.new_stop_iteration() - })?; + let ch = match value[start..].chars().next() { + Some(ch) => ch, + None => { + zelf.status.store(Exhausted); + return Ok(PyIterReturn::StopIteration(None)); + } + }; match zelf.position.compare_exchange_weak( start, @@ -252,7 +256,7 @@ impl SlotIterator for PyStrIterator { atomic::Ordering::Release, atomic::Ordering::Relaxed, ) { - Ok(_) => break Ok(ch.into_pyobject(vm)), + Ok(_) => break Ok(PyIterReturn::Return(ch.into_pyobject(vm))), Err(cur) => start = cur, } } diff --git a/vm/src/builtins/pytype.rs b/vm/src/builtins/pytype.rs index 71a27a327..b9f67d096 100644 --- a/vm/src/builtins/pytype.rs +++ b/vm/src/builtins/pytype.rs @@ -5,6 +5,7 @@ use super::{ use crate::common::{ascii, lock::PyRwLock}; use crate::{ function::{FuncArgs, KwArgs, OptionalArg}, + protocol::PyIterReturn, slots::{self, Callable, PyTypeFlags, PyTypeSlots, SlotGetattro, SlotSetattro}, utils::Either, IdProtocol, PyAttributes, PyClassImpl, PyContext, PyLease, PyObjectRef, PyRef, PyResult, @@ -261,8 +262,12 @@ impl PyType { update_slot!(iter, func); } "__next__" => { - let func: slots::IterNextFunc = - |zelf, vm| vm.call_special_method(zelf.clone(), "__next__", ()); + let func: slots::IterNextFunc = |zelf, vm| { + PyIterReturn::from_result( + vm.call_special_method(zelf.clone(), "__next__", ()), + vm, + ) + }; update_slot!(iternext, func); } _ => {} diff --git a/vm/src/builtins/range.rs b/vm/src/builtins/range.rs index bfb371a4b..4f6db0a59 100644 --- a/vm/src/builtins/range.rs +++ b/vm/src/builtins/range.rs @@ -2,7 +2,7 @@ use super::{PyInt, PyIntRef, PySlice, PySliceRef, PyTypeRef}; use crate::common::hash::PyHash; use crate::{ function::{FuncArgs, OptionalArg}, - iterator, + protocol::PyIterReturn, slots::{Comparable, Hashable, Iterable, IteratorIterable, PyComparisonOp, SlotIterator}, IdProtocol, IntoPyRef, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol, VirtualMachine, @@ -31,7 +31,7 @@ fn iter_search( ) -> PyResult { let mut count = 0; let iter = obj.get_iter(vm)?; - while let Some(element) = iterator::get_next_object(vm, &iter)? { + while let PyIterReturn::Return(element) = iter.next(vm)? { if vm.bool_eq(&item, &element)? { match flag { SearchType::Index => return Ok(count), @@ -515,17 +515,18 @@ impl PyLongRangeIterator { impl IteratorIterable for PyLongRangeIterator {} impl SlotIterator for PyLongRangeIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { // TODO: In pathological case (index == usize::MAX) this can wrap around // (since fetch_add wraps). This would result in the iterator spinning again // from the beginning. let index = BigInt::from(zelf.index.fetch_add(1)); if index < zelf.length { - Ok(vm - .ctx - .new_int(zelf.start.clone() + index * zelf.step.clone())) + Ok(PyIterReturn::Return( + vm.ctx + .new_int(zelf.start.clone() + index * zelf.step.clone()), + )) } else { - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } @@ -585,15 +586,17 @@ impl PyRangeIterator { impl IteratorIterable for PyRangeIterator {} impl SlotIterator for PyRangeIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { // TODO: In pathological case (index == usize::MAX) this can wrap around // (since fetch_add wraps). This would result in the iterator spinning again // from the beginning. let index = zelf.index.fetch_add(1); if index < zelf.length { - Ok(vm.ctx.new_int(zelf.start + (index as isize) * zelf.step)) + Ok(PyIterReturn::Return( + vm.ctx.new_int(zelf.start + (index as isize) * zelf.step), + )) } else { - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } diff --git a/vm/src/builtins/set.rs b/vm/src/builtins/set.rs index 7a6510b22..f162f5891 100644 --- a/vm/src/builtins/set.rs +++ b/vm/src/builtins/set.rs @@ -6,6 +6,7 @@ use crate::common::{ascii, hash::PyHash, rc::PyRc}; use crate::{ dictdatatype::{self, DictSize}, function::{ArgIterable, FuncArgs, OptionalArg, PosArgs}, + protocol::PyIterReturn, slots::{ Comparable, Hashable, Iterable, IteratorIterable, PyComparisonOp, SlotConstructor, SlotIterator, Unhashable, @@ -863,9 +864,9 @@ impl PySetIterator { impl IteratorIterable for PySetIterator {} impl SlotIterator for PySetIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { match zelf.status.load() { - IterStatus::Exhausted => Err(vm.new_stop_iteration()), + IterStatus::Exhausted => Ok(PyIterReturn::StopIteration(None)), IterStatus::Active => { if zelf.dict.has_changed_size(&zelf.size) { zelf.status.store(IterStatus::Exhausted); @@ -874,10 +875,10 @@ impl SlotIterator for PySetIterator { ); } match zelf.dict.next_entry_atomic(&zelf.position) { - Some((key, _)) => Ok(key), + Some((key, _)) => Ok(PyIterReturn::Return(key)), None => { zelf.status.store(IterStatus::Exhausted); - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } diff --git a/vm/src/builtins/tuple.rs b/vm/src/builtins/tuple.rs index 65666cbf2..f733f222a 100644 --- a/vm/src/builtins/tuple.rs +++ b/vm/src/builtins/tuple.rs @@ -6,6 +6,7 @@ use super::{ use crate::common::hash::PyHash; use crate::{ function::OptionalArg, + protocol::PyIterReturn, sequence::{self, SimpleSeq}, sliceable::PySliceableSequence, slots::{ @@ -378,16 +379,16 @@ impl PyTupleIterator { impl IteratorIterable for PyTupleIterator {} impl SlotIterator for PyTupleIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, _vm: &VirtualMachine) -> PyResult { if let Exhausted = zelf.status.load() { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } let pos = zelf.position.fetch_add(1); if let Some(obj) = zelf.tuple.as_slice().get(pos) { - Ok(obj.clone()) + Ok(PyIterReturn::Return(obj.clone())) } else { zelf.status.store(Exhausted); - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } diff --git a/vm/src/builtins/zip.rs b/vm/src/builtins/zip.rs index 8412803d9..c5a2683ee 100644 --- a/vm/src/builtins/zip.rs +++ b/vm/src/builtins/zip.rs @@ -1,7 +1,7 @@ use super::PyTypeRef; use crate::{ function::PosArgs, - protocol::PyIter, + protocol::{PyIter, PyIterReturn}, slots::{IteratorIterable, SlotConstructor, SlotIterator}, PyClassImpl, PyContext, PyRef, PyResult, PyValue, VirtualMachine, }; @@ -32,18 +32,19 @@ impl PyZip {} impl IteratorIterable for PyZip {} impl SlotIterator for PyZip { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { if zelf.iterators.is_empty() { - Err(vm.new_stop_iteration()) - } else { - let next_objs = zelf - .iterators - .iter() - .map(|iterator| iterator.next(vm)) - .collect::, _>>()?; - - Ok(vm.ctx.new_tuple(next_objs)) + return Ok(PyIterReturn::StopIteration(None)); } + let mut next_objs = Vec::new(); + for iterator in zelf.iterators.iter() { + let item = match iterator.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + }; + next_objs.push(item); + } + Ok(PyIterReturn::Return(vm.ctx.new_tuple(next_objs))) } } diff --git a/vm/src/frame.rs b/vm/src/frame.rs index 2d77e3eae..02d345dcd 100644 --- a/vm/src/frame.rs +++ b/vm/src/frame.rs @@ -15,12 +15,12 @@ use crate::{ exceptions::{self, ExceptionCtor}, function::FuncArgs, iterator, - protocol::PyIter, + protocol::{PyIter, PyIterReturn}, scope::Scope, slots::PyComparisonOp, stdlib::builtins, - IdProtocol, ItemProtocol, PyMethod, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, - TypeProtocol, VirtualMachine, + IdProtocol, IntoPyResult, ItemProtocol, PyMethod, PyObjectRef, PyRef, PyResult, PyValue, + TryFromObject, TypeProtocol, VirtualMachine, }; use indexmap::IndexMap; use itertools::Itertools; @@ -1442,7 +1442,8 @@ impl ExecutingFrame<'_> { fn _send(&self, coro: &PyObjectRef, val: PyObjectRef, vm: &VirtualMachine) -> PyResult { match self.builtin_coro(coro) { Some(coro) => coro.send(val, vm), - None if vm.is_none(&val) => PyIter::new(coro).next(vm), + // FIXME: turn return type to PyResult then ExecutionResult will be simplified + None if vm.is_none(&val) => PyIter::new(coro).next(vm).into_pyresult(vm), None => { let meth = vm.get_attribute(coro.clone(), "send")?; vm.invoke(&meth, (val,)) @@ -1513,15 +1514,15 @@ impl ExecutingFrame<'_> { /// The top of stack contains the iterator, lets push it forward fn execute_for_iter(&mut self, vm: &VirtualMachine, target: bytecode::Label) -> FrameResult { let top_of_stack = PyIter::new(self.last_value()); - let next_obj = iterator::get_next_object(vm, &top_of_stack); + let next_obj = top_of_stack.next(vm); // Check the next object: match next_obj { - Ok(Some(value)) => { + Ok(PyIterReturn::Return(value)) => { self.push_value(value); Ok(None) } - Ok(None) => { + Ok(PyIterReturn::StopIteration(_)) => { // Pop iterator from stack: self.pop_value(); diff --git a/vm/src/function/argument.rs b/vm/src/function/argument.rs index bc87e1dbf..d60a65c13 100644 --- a/vm/src/function/argument.rs +++ b/vm/src/function/argument.rs @@ -1,7 +1,9 @@ use super::IntoFuncArgs; use crate::{ - builtins::iter::PySequenceIterator, iterator, protocol::PyIter, PyObjectRef, PyResult, PyValue, - TryFromObject, TypeProtocol, VirtualMachine, + builtins::iter::PySequenceIterator, + iterator, + protocol::{PyIter, PyIterReturn}, + PyObjectRef, PyResult, PyValue, TryFromObject, TypeProtocol, VirtualMachine, }; use std::marker::PhantomData; @@ -102,7 +104,12 @@ where type Item = PyResult; fn next(&mut self) -> Option { - iterator::get_next_object(self.vm, &PyIter::new(&self.obj)) + PyIter::new(&self.obj) + .next(self.vm) + .map(|iret| match iret { + PyIterReturn::Return(obj) => Some(obj), + PyIterReturn::StopIteration(_) => None, + }) .transpose() .map(|x| x.and_then(|obj| T::try_from_object(self.vm, obj))) } diff --git a/vm/src/iterator.rs b/vm/src/iterator.rs index 3bbb0d043..3509ae410 100644 --- a/vm/src/iterator.rs +++ b/vm/src/iterator.rs @@ -4,36 +4,11 @@ use crate::{ builtins::{int, PyBaseExceptionRef, PyInt}, - protocol::PyIter, + protocol::{PyIter, PyIterReturn}, IdProtocol, PyObjectRef, PyResult, TypeProtocol, VirtualMachine, }; use num_traits::Signed; -/* - * Helper function to retrieve the next object (or none) from an iterator. - */ -pub fn get_next_object( - vm: &VirtualMachine, - iter_obj: &PyIter, -) -> PyResult> -where - T: std::borrow::Borrow, -{ - let next_obj: PyResult = iter_obj.next(vm); - - match next_obj { - Ok(value) => Ok(Some(value)), - Err(next_error) => { - // Check if we have stopiteration, or something else: - if next_error.isinstance(&vm.ctx.exceptions.stop_iteration) { - Ok(None) - } else { - Err(next_error) - } - } - } -} - pub fn try_map(vm: &VirtualMachine, iter: &PyIter, cap: usize, mut f: F) -> PyResult> where F: FnMut(PyObjectRef) -> PyResult, @@ -44,7 +19,7 @@ where return Ok(Vec::new()); } let mut results = Vec::with_capacity(cap); - while let Some(element) = get_next_object(vm, iter)? { + while let PyIterReturn::Return(element) = iter.next(vm)? { results.push(f(element)?); } results.shrink_to_fit(); diff --git a/vm/src/protocol/iter.rs b/vm/src/protocol/iter.rs index 999e7baa5..967333bff 100644 --- a/vm/src/protocol/iter.rs +++ b/vm/src/protocol/iter.rs @@ -1,7 +1,7 @@ use crate::IntoPyObject; use crate::{ - builtins::iter::PySequenceIterator, PyObjectRef, PyResult, PyValue, TryFromObject, - TypeProtocol, VirtualMachine, + builtins::iter::PySequenceIterator, IntoPyResult, PyObjectRef, PyResult, PyValue, + TryFromObject, TypeProtocol, VirtualMachine, }; use std::borrow::Borrow; use std::ops::Deref; @@ -35,7 +35,7 @@ where pub fn as_object(&self) -> &PyObjectRef { self.0.borrow() } - pub fn next(&self, vm: &VirtualMachine) -> PyResult { + pub fn next(&self, vm: &VirtualMachine) -> PyResult { let iternext = { self.0 .borrow() @@ -119,3 +119,39 @@ impl PyObjectRef { PyIter::try_from_object(vm, self) } } + +pub enum PyIterReturn { + Return(T), + StopIteration(Option), +} + +impl PyIterReturn { + pub fn from_result(result: PyResult, vm: &VirtualMachine) -> PyResult { + match result { + Ok(obj) => Ok(Self::Return(obj)), + Err(err) if err.isinstance(&vm.ctx.exceptions.stop_iteration) => { + let args = err.get_arg(0); + Ok(Self::StopIteration(args)) + } + Err(err) => Err(err), + } + } +} + +impl IntoPyResult for PyIterReturn { + fn into_pyresult(self, vm: &VirtualMachine) -> PyResult { + match self { + Self::Return(obj) => Ok(obj), + Self::StopIteration(v) => Err({ + let args = if let Some(v) = v { vec![v] } else { Vec::new() }; + vm.new_exception(vm.ctx.exceptions.stop_iteration.clone(), args) + }), + } + } +} + +impl IntoPyResult for PyResult { + fn into_pyresult(self, vm: &VirtualMachine) -> PyResult { + self.and_then(|obj| obj.into_pyresult(vm)) + } +} diff --git a/vm/src/protocol/mod.rs b/vm/src/protocol/mod.rs index 24692be4c..0c44ab754 100644 --- a/vm/src/protocol/mod.rs +++ b/vm/src/protocol/mod.rs @@ -2,4 +2,4 @@ mod buffer; mod iter; pub use buffer::{BufferInternal, BufferOptions, PyBuffer, ResizeGuard}; -pub use iter::PyIter; +pub use iter::{PyIter, PyIterReturn}; diff --git a/vm/src/slots.rs b/vm/src/slots.rs index 4bf793362..37382e6ec 100644 --- a/vm/src/slots.rs +++ b/vm/src/slots.rs @@ -2,12 +2,12 @@ use crate::builtins::{PyStrRef, PyTypeRef}; use crate::common::hash::PyHash; use crate::common::lock::PyRwLock; use crate::function::{FromArgs, FuncArgs, OptionalArg}; -use crate::protocol::PyBuffer; +use crate::protocol::{PyBuffer, PyIterReturn}; use crate::utils::Either; use crate::VirtualMachine; use crate::{ - IdProtocol, PyComparisonValue, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, - TypeProtocol, + IdProtocol, IntoPyResult, PyComparisonValue, PyObjectRef, PyRef, PyResult, PyValue, + TryFromObject, TypeProtocol, }; use crossbeam_utils::atomic::AtomicCell; use std::cmp::Ordering; @@ -71,7 +71,7 @@ pub(crate) type SetattroFunc = fn(&PyObjectRef, PyStrRef, Option, &VirtualMachine) -> PyResult<()>; pub(crate) type BufferFunc = fn(&PyObjectRef, &VirtualMachine) -> PyResult; pub(crate) type IterFunc = fn(PyObjectRef, &VirtualMachine) -> PyResult; -pub(crate) type IterNextFunc = fn(&PyObjectRef, &VirtualMachine) -> PyResult; +pub(crate) type IterNextFunc = fn(&PyObjectRef, &VirtualMachine) -> PyResult; // The corresponding field in CPython is `tp_` prefixed. // e.g. name -> tp_name @@ -550,7 +550,7 @@ pub trait Iterable: PyValue { #[pyimpl(with(Iterable))] pub trait SlotIterator: PyValue + Iterable { #[pyslot] - fn slot_iternext(zelf: &PyObjectRef, vm: &VirtualMachine) -> PyResult { + fn slot_iternext(zelf: &PyObjectRef, vm: &VirtualMachine) -> PyResult { if let Some(zelf) = zelf.downcast_ref() { Self::next(zelf, vm) } else { @@ -558,11 +558,11 @@ pub trait SlotIterator: PyValue + Iterable { } } - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult; + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult; #[pymethod] fn __next__(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult { - Self::slot_iternext(&zelf, vm) + Self::slot_iternext(&zelf, vm).into_pyresult(vm) } } diff --git a/vm/src/stdlib/builtins.rs b/vm/src/stdlib/builtins.rs index ab4c6e591..879bf4c68 100644 --- a/vm/src/stdlib/builtins.rs +++ b/vm/src/stdlib/builtins.rs @@ -25,7 +25,7 @@ mod builtins { ArgBytesLike, ArgCallable, ArgIterable, FuncArgs, KwArgs, OptionalArg, OptionalOption, PosArgs, }, - protocol::PyIter, + protocol::{PyIter, PyIterReturn}, py_io, readline::{Readline, ReadlineResult}, scope::Scope, @@ -512,19 +512,17 @@ mod builtins { iterator: PyObjectRef, default_value: OptionalArg, vm: &VirtualMachine, - ) -> PyResult { + ) -> PyResult { if !PyIter::check(&iterator) { return Err(vm.new_type_error(format!( "{} object is not an iterator", iterator.class().name() ))); } - - PyIter::new(iterator).next(vm).or_else(|err| { - if err.isinstance(&vm.ctx.exceptions.stop_iteration) { - default_value.ok_or(err) - } else { - Err(err) + PyIter::new(iterator).next(vm).map(|iret| match iret { + PyIterReturn::Return(obj) => PyIterReturn::Return(obj), + PyIterReturn::StopIteration(v) => { + default_value.map_or(PyIterReturn::StopIteration(v), PyIterReturn::Return) } }) } diff --git a/vm/src/stdlib/collections.rs b/vm/src/stdlib/collections.rs index 4a24b3219..9158ea67a 100644 --- a/vm/src/stdlib/collections.rs +++ b/vm/src/stdlib/collections.rs @@ -9,6 +9,7 @@ mod _collections { PyInt, PyTypeRef, }, function::{FuncArgs, KwArgs, OptionalArg}, + protocol::PyIterReturn, sequence, sliceable, slots::{ Comparable, Hashable, Iterable, IteratorIterable, PyComparisonOp, SlotConstructor, @@ -642,9 +643,9 @@ mod _collections { impl IteratorIterable for PyDequeIterator {} impl SlotIterator for PyDequeIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { match zelf.status.load() { - Exhausted => Err(vm.new_stop_iteration()), + Exhausted => Ok(PyIterReturn::StopIteration(None)), Active => { if zelf.length != zelf.deque.len() { // Deque was changed while we iterated. @@ -655,10 +656,10 @@ mod _collections { let deque = zelf.deque.borrow_deque(); if pos < deque.len() { let ret = deque[pos].clone(); - Ok(ret) + Ok(PyIterReturn::Return(ret)) } else { zelf.status.store(Exhausted); - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } @@ -722,9 +723,9 @@ mod _collections { impl IteratorIterable for PyReverseDequeIterator {} impl SlotIterator for PyReverseDequeIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { match zelf.status.load() { - Exhausted => Err(vm.new_stop_iteration()), + Exhausted => Ok(PyIterReturn::StopIteration(None)), Active => { // If length changes while we iterate, set to Exhausted and bail. if zelf.length != zelf.deque.len() { @@ -735,7 +736,7 @@ mod _collections { let deque = zelf.deque.borrow_deque(); if pos > 0 { if let Some(obj) = deque.get(pos) { - return Ok(obj.clone()); + return Ok(PyIterReturn::Return(obj.clone())); } } // We either are == 0 or deque.get returned None. Either way, set status @@ -743,9 +744,9 @@ mod _collections { zelf.status.store(Exhausted); if pos == 0 { // Can safely index directly. - return Ok(deque[pos].clone()); + return Ok(PyIterReturn::Return(deque[pos].clone())); } - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } diff --git a/vm/src/stdlib/functools.rs b/vm/src/stdlib/functools.rs index 08047e792..f3eb94ebc 100644 --- a/vm/src/stdlib/functools.rs +++ b/vm/src/stdlib/functools.rs @@ -3,9 +3,9 @@ pub(crate) use _functools::make_module; #[pymodule] mod _functools { use crate::function::OptionalArg; - use crate::protocol::PyIter; + use crate::protocol::{PyIter, PyIterReturn}; use crate::vm::VirtualMachine; - use crate::{PyObjectRef, PyResult, TypeProtocol}; + use crate::{PyObjectRef, PyResult}; #[pyfunction] fn reduce( @@ -17,22 +17,25 @@ mod _functools { let start_value = if let OptionalArg::Present(val) = start_value { val } else { - iterator.next(vm).map_err(|err| { - if err.isinstance(&vm.ctx.exceptions.stop_iteration) { + iterator.next(vm).and_then(|iret| match iret { + PyIterReturn::Return(obj) => Ok(obj), + PyIterReturn::StopIteration(_) => Err({ let exc_type = vm.ctx.exceptions.type_error.clone(); vm.new_exception_msg( exc_type, "reduce() of empty sequence with no initial value".to_owned(), ) - } else { - err - } + }), })? }; let mut accumulator = start_value; while let Ok(next_obj) = iterator.next(vm) { + let next_obj = match next_obj { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(_) => break, + }; accumulator = vm.invoke(&function, vec![accumulator, next_obj])? } diff --git a/vm/src/stdlib/io.rs b/vm/src/stdlib/io.rs index ba705f704..eee69a9a1 100644 --- a/vm/src/stdlib/io.rs +++ b/vm/src/stdlib/io.rs @@ -85,7 +85,7 @@ mod _io { function::{ ArgBytesLike, ArgIterable, ArgMemoryBuffer, FuncArgs, OptionalArg, OptionalOption, }, - protocol::{BufferInternal, BufferOptions, PyBuffer, ResizeGuard}, + protocol::{BufferInternal, BufferOptions, PyBuffer, PyIterReturn, ResizeGuard}, slots::{Iterable, SlotConstructor, SlotIterator}, utils::Either, vm::{ReprGuard, VirtualMachine}, @@ -524,16 +524,16 @@ mod _io { } impl SlotIterator for _IOBase { - fn slot_iternext(zelf: &PyObjectRef, vm: &VirtualMachine) -> PyResult { + fn slot_iternext(zelf: &PyObjectRef, vm: &VirtualMachine) -> PyResult { let line = vm.call_method(zelf, "readline", ())?; - if !line.clone().try_to_bool(vm)? { - Err(vm.new_stop_iteration()) + Ok(if !line.clone().try_to_bool(vm)? { + PyIterReturn::StopIteration(None) } else { - Ok(line) - } + PyIterReturn::Return(line) + }) } - fn next(_zelf: &PyRef, _vm: &VirtualMachine) -> PyResult { + fn next(_zelf: &PyRef, _vm: &VirtualMachine) -> PyResult { unreachable!("slot_iternext is implemented") } } diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index d0dcb1274..d81227a0a 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -9,8 +9,7 @@ mod decl { use crate::{ builtins::{int, PyInt, PyIntRef, PyTupleRef, PyTypeRef}, function::{ArgCallable, FuncArgs, OptionalArg, OptionalOption, PosArgs}, - iterator::get_next_object, - protocol::PyIter, + protocol::{PyIter, PyIterReturn}, slots::{IteratorIterable, SlotConstructor, SlotIterator}, IdProtocol, IntoPyObject, PyObjectRef, PyRef, PyResult, PyValue, PyWeakRef, TypeProtocol, VirtualMachine, @@ -57,7 +56,7 @@ mod decl { } impl IteratorIterable for PyItertoolsChain {} impl SlotIterator for PyItertoolsChain { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { loop { let pos = zelf.cur_idx.load(); if pos >= zelf.iterables.len() { @@ -77,19 +76,18 @@ mod decl { // We need to call "next" outside of the lock. match cur_iter.next(vm) { - Ok(ok) => return Ok(ok), + Ok(PyIterReturn::Return(ok)) => return Ok(PyIterReturn::Return(ok)), + Ok(PyIterReturn::StopIteration(_)) => { + zelf.cur_idx.fetch_add(1); + *zelf.cached_iter.write() = None; + } Err(err) => { - if err.isinstance(&vm.ctx.exceptions.stop_iteration) { - zelf.cur_idx.fetch_add(1); - *zelf.cached_iter.write() = None; - } else { - return Err(err); - } + return Err(err); } } } - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } @@ -126,9 +124,12 @@ mod decl { impl IteratorIterable for PyItertoolsCompress {} impl SlotIterator for PyItertoolsCompress { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { loop { - let sel_obj = zelf.selector.next(vm)?; + let sel_obj = match zelf.selector.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + }; let verdict = sel_obj.clone().try_to_bool(vm)?; let data_obj = zelf.data.next(vm)?; @@ -185,11 +186,11 @@ mod decl { impl PyItertoolsCount {} impl IteratorIterable for PyItertoolsCount {} impl SlotIterator for PyItertoolsCount { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let mut cur = zelf.cur.write(); let result = cur.clone(); *cur += &zelf.step; - Ok(result.into_pyobject(vm)) + Ok(PyIterReturn::Return(result.into_pyobject(vm))) } } @@ -219,14 +220,14 @@ mod decl { impl PyItertoolsCycle {} impl IteratorIterable for PyItertoolsCycle {} impl SlotIterator for PyItertoolsCycle { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - let item = if let Some(item) = get_next_object(vm, &zelf.iter)? { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let item = if let PyIterReturn::Return(item) = zelf.iter.next(vm)? { zelf.saved.write().push(item.clone()); item } else { let saved = zelf.saved.read(); if saved.len() == 0 { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } let last_index = zelf.index.fetch_add(1); @@ -238,7 +239,7 @@ mod decl { saved[last_index].clone() }; - Ok(item) + Ok(PyIterReturn::Return(item)) } } @@ -320,15 +321,15 @@ mod decl { impl IteratorIterable for PyItertoolsRepeat {} impl SlotIterator for PyItertoolsRepeat { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, _vm: &VirtualMachine) -> PyResult { if let Some(ref times) = zelf.times { let mut times = times.write(); if *times == 0 { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } *times -= 1; } - Ok(zelf.object.clone()) + Ok(PyIterReturn::Return(zelf.object.clone())) } } @@ -365,11 +366,15 @@ mod decl { impl PyItertoolsStarmap {} impl IteratorIterable for PyItertoolsStarmap {} impl SlotIterator for PyItertoolsStarmap { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let obj = zelf.iter.next(vm)?; let function = &zelf.function; - - vm.invoke(function, vm.extract_elements(&obj)?) + match obj { + PyIterReturn::Return(obj) => { + PyIterReturn::from_result(vm.invoke(function, vm.extract_elements(&obj)?), vm) + } + PyIterReturn::StopIteration(v) => Ok(PyIterReturn::StopIteration(v)), + } } } @@ -414,22 +419,25 @@ mod decl { impl PyItertoolsTakewhile {} impl IteratorIterable for PyItertoolsTakewhile {} impl SlotIterator for PyItertoolsTakewhile { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { if zelf.stop_flag.load() { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } // might be StopIteration or anything else, which is propagated upwards - let obj = zelf.iterable.next(vm)?; + let obj = match zelf.iterable.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + }; let predicate = &zelf.predicate; let verdict = vm.invoke(predicate, (obj.clone(),))?; let verdict = verdict.try_to_bool(vm)?; if verdict { - Ok(obj) + Ok(PyIterReturn::Return(obj)) } else { zelf.stop_flag.store(true); - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } @@ -475,18 +483,23 @@ mod decl { impl PyItertoolsDropwhile {} impl IteratorIterable for PyItertoolsDropwhile {} impl SlotIterator for PyItertoolsDropwhile { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let predicate = &zelf.predicate; let iterable = &zelf.iterable; if !zelf.start_flag.load() { loop { - let obj = iterable.next(vm)?; + let obj = match iterable.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => { + return Ok(PyIterReturn::StopIteration(v)) + } + }; let pred = predicate.clone(); let pred_value = vm.invoke(&pred.into_object(), (obj.clone(),))?; if !pred_value.try_to_bool(vm)? { zelf.start_flag.store(true); - return Ok(obj); + return Ok(PyIterReturn::Return(obj)); } } } @@ -570,19 +583,25 @@ mod decl { #[pyimpl(with(SlotIterator, SlotConstructor))] impl PyItertoolsGroupBy { - pub(super) fn advance(&self, vm: &VirtualMachine) -> PyResult<(PyObjectRef, PyObjectRef)> { - let new_value = self.iterable.next(vm)?; + pub(super) fn advance( + &self, + vm: &VirtualMachine, + ) -> PyResult> { + let new_value = match self.iterable.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + }; let new_key = if let Some(ref kf) = self.key_func { vm.invoke(kf, vec![new_value.clone()])? } else { new_value.clone() }; - Ok((new_value, new_key)) + Ok(PyIterReturn::Return((new_value, new_key))) } } impl IteratorIterable for PyItertoolsGroupBy {} impl SlotIterator for PyItertoolsGroupBy { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let mut state = zelf.state.lock(); state.grouper = None; @@ -593,13 +612,23 @@ mod decl { let (value, key) = if let Some(old_key) = current_key { loop { - let (value, new_key) = zelf.advance(vm)?; + let (value, new_key) = match zelf.advance(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => { + return Ok(PyIterReturn::StopIteration(v)) + } + }; if !vm.bool_eq(&new_key, &old_key)? { break (value, new_key); } } } else { - zelf.advance(vm)? + match zelf.advance(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => { + return Ok(PyIterReturn::StopIteration(v)) + } + } }; state = zelf.state.lock(); @@ -615,7 +644,9 @@ mod decl { .into_ref(vm); state.grouper = Some(PyRef::downgrade(&grouper)); - Ok((state.current_key.as_ref().unwrap().clone(), grouper).into_pyobject(vm)) + Ok(PyIterReturn::Return( + (state.current_key.as_ref().unwrap().clone(), grouper).into_pyobject(vm), + )) } } @@ -632,31 +663,34 @@ mod decl { impl PyItertoolsGrouper {} impl IteratorIterable for PyItertoolsGrouper {} impl SlotIterator for PyItertoolsGrouper { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let old_key = { let mut state = zelf.groupby.state.lock(); if !state.is_current(zelf) { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } // check to see if the value has already been retrieved from the iterator if let Some(val) = state.current_value.take() { - return Ok(val); + return Ok(PyIterReturn::Return(val)); } state.current_key.as_ref().unwrap().clone() }; - let (value, key) = zelf.groupby.advance(vm)?; + let (value, key) = match zelf.groupby.advance(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + }; if vm.bool_eq(&key, &old_key)? { - Ok(value) + Ok(PyIterReturn::Return(value)) } else { let mut state = zelf.groupby.state.lock(); state.current_value = Some(value); state.current_key = Some(key); state.next_group = true; state.grouper = None; - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } @@ -762,7 +796,7 @@ mod decl { impl IteratorIterable for PyItertoolsIslice {} impl SlotIterator for PyItertoolsIslice { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { while zelf.cur.load() < zelf.next.load() { zelf.iterable.next(vm)?; zelf.cur.fetch_add(1); @@ -770,18 +804,21 @@ mod decl { if let Some(stop) = zelf.stop { if zelf.cur.load() >= stop { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } } - let obj = zelf.iterable.next(vm)?; + let obj = match zelf.iterable.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + }; zelf.cur.fetch_add(1); // TODO is this overflow check required? attempts to copy CPython. let (next, ovf) = zelf.next.load().overflowing_add(zelf.step); zelf.next.store(if ovf { zelf.stop.unwrap() } else { next }); - Ok(obj) + Ok(PyIterReturn::Return(obj)) } } @@ -824,12 +861,15 @@ mod decl { impl PyItertoolsFilterFalse {} impl IteratorIterable for PyItertoolsFilterFalse {} impl SlotIterator for PyItertoolsFilterFalse { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let predicate = &zelf.predicate; let iterable = &zelf.iterable; loop { - let obj = iterable.next(vm)?; + let obj = match iterable.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + }; let pred_value = if vm.is_none(predicate) { obj.clone() } else { @@ -837,7 +877,7 @@ mod decl { }; if !pred_value.try_to_bool(vm)? { - return Ok(obj); + return Ok(PyIterReturn::Return(obj)); } } } @@ -881,18 +921,28 @@ mod decl { impl IteratorIterable for PyItertoolsAccumulate {} impl SlotIterator for PyItertoolsAccumulate { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let iterable = &zelf.iterable; let acc_value = zelf.acc_value.read().clone(); let next_acc_value = match acc_value { None => match &zelf.initial { - None => iterable.next(vm)?, + None => match iterable.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => { + return Ok(PyIterReturn::StopIteration(v)) + } + }, Some(obj) => obj.clone(), }, Some(value) => { - let obj = iterable.next(vm)?; + let obj = match iterable.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => { + return Ok(PyIterReturn::StopIteration(v)) + } + }; match &zelf.binop { None => vm._add(&value, &obj)?, Some(op) => vm.invoke(op, vec![value, obj])?, @@ -901,7 +951,7 @@ mod decl { }; *zelf.acc_value.write() = Some(next_acc_value.clone()); - Ok(next_acc_value) + Ok(PyIterReturn::Return(next_acc_value)) } } @@ -919,12 +969,15 @@ mod decl { })) } - fn get_item(&self, vm: &VirtualMachine, index: usize) -> PyResult { + fn get_item(&self, vm: &VirtualMachine, index: usize) -> PyResult { if self.values.read().len() == index { - let result = self.iterable.next(vm)?; + let result = match self.iterable.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + }; self.values.write().push(result); } - Ok(self.values.read()[index].clone()) + Ok(PyIterReturn::Return(self.values.read()[index].clone())) } } @@ -999,10 +1052,13 @@ mod decl { } impl IteratorIterable for PyItertoolsTee {} impl SlotIterator for PyItertoolsTee { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - let value = zelf.tee_data.get_item(vm, zelf.index.load())?; + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let value = match zelf.tee_data.get_item(vm, zelf.index.load())? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + }; zelf.index.fetch_add(1); - Ok(value) + Ok(PyIterReturn::Return(value)) } } @@ -1075,17 +1131,17 @@ mod decl { } impl IteratorIterable for PyItertoolsProduct {} impl SlotIterator for PyItertoolsProduct { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { // stop signal if zelf.stop.load() { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } let pools = &zelf.pools; for p in pools { if p.is_empty() { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } } @@ -1100,7 +1156,7 @@ mod decl { zelf.update_idxs(idxs); - Ok(res) + Ok(PyIterReturn::Return(res)) } } @@ -1154,10 +1210,10 @@ mod decl { impl PyItertoolsCombinations {} impl IteratorIterable for PyItertoolsCombinations {} impl SlotIterator for PyItertoolsCombinations { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { // stop signal if zelf.exhausted.load() { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } let n = zelf.pool.len(); @@ -1165,7 +1221,7 @@ mod decl { if r == 0 { zelf.exhausted.store(true); - return Ok(vm.ctx.new_tuple(vec![])); + return Ok(PyIterReturn::Return(vm.ctx.new_tuple(vec![]))); } let res = vm.ctx.new_tuple( @@ -1199,7 +1255,7 @@ mod decl { } } - Ok(res) + Ok(PyIterReturn::Return(res)) } } @@ -1245,10 +1301,10 @@ mod decl { impl IteratorIterable for PyItertoolsCombinationsWithReplacement {} impl SlotIterator for PyItertoolsCombinationsWithReplacement { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { // stop signal if zelf.exhausted.load() { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } let n = zelf.pool.len(); @@ -1256,7 +1312,7 @@ mod decl { if r == 0 { zelf.exhausted.store(true); - return Ok(vm.ctx.new_tuple(vec![])); + return Ok(PyIterReturn::Return(vm.ctx.new_tuple(vec![]))); } let mut indices = zelf.indices.write(); @@ -1285,7 +1341,7 @@ mod decl { } } - Ok(res) + Ok(PyIterReturn::Return(res)) } } @@ -1353,10 +1409,10 @@ mod decl { impl PyItertoolsPermutations {} impl IteratorIterable for PyItertoolsPermutations {} impl SlotIterator for PyItertoolsPermutations { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { // stop signal if zelf.exhausted.load() { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } let n = zelf.pool.len(); @@ -1364,7 +1420,7 @@ mod decl { if n == 0 { zelf.exhausted.store(true); - return Ok(vm.ctx.new_tuple(vec![])); + return Ok(PyIterReturn::Return(vm.ctx.new_tuple(vec![]))); } let mut result = zelf.result.write(); @@ -1401,20 +1457,22 @@ mod decl { } if !sentinel { zelf.exhausted.store(true); - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } } else { // On the first pass, initialize result tuple using the indices *result = Some((0..r).collect()); } - Ok(vm.ctx.new_tuple( - result - .as_ref() - .unwrap() - .iter() - .map(|&i| zelf.pool[i].clone()) - .collect(), + Ok(PyIterReturn::Return( + vm.ctx.new_tuple( + result + .as_ref() + .unwrap() + .iter() + .map(|&i| zelf.pool[i].clone()) + .collect(), + ), )) } } @@ -1451,30 +1509,27 @@ mod decl { impl PyItertoolsZipLongest {} impl IteratorIterable for PyItertoolsZipLongest {} impl SlotIterator for PyItertoolsZipLongest { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { if zelf.iterators.is_empty() { - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } else { let mut result: Vec = Vec::new(); let mut numactive = zelf.iterators.len(); for idx in 0..zelf.iterators.len() { - let next_obj = match zelf.iterators[idx].next(vm) { - Ok(obj) => obj, - Err(err) => { - if !err.isinstance(&vm.ctx.exceptions.stop_iteration) { - return Err(err); - } + let next_obj = match zelf.iterators[idx].next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => { numactive -= 1; if numactive == 0 { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(v)); } zelf.fillvalue.clone() } }; result.push(next_obj); } - Ok(vm.ctx.new_tuple(result)) + Ok(PyIterReturn::Return(vm.ctx.new_tuple(result))) } } } @@ -1503,14 +1558,20 @@ mod decl { impl PyItertoolsPairwise {} impl IteratorIterable for PyItertoolsPairwise {} impl SlotIterator for PyItertoolsPairwise { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let old = match zelf.old.read().clone() { - None => zelf.iterator.next(vm)?, + None => match zelf.iterator.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + }, Some(obj) => obj, }; - let new = zelf.iterator.next(vm)?; + let new = match zelf.iterator.next(vm)? { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + }; *zelf.old.write() = Some(new.clone()); - Ok(vm.ctx.new_tuple(vec![old, new])) + Ok(PyIterReturn::Return(vm.ctx.new_tuple(vec![old, new]))) } } } diff --git a/vm/src/stdlib/operator.rs b/vm/src/stdlib/operator.rs index d71359b8a..748c0a7de 100644 --- a/vm/src/stdlib/operator.rs +++ b/vm/src/stdlib/operator.rs @@ -14,7 +14,7 @@ mod _operator { builtins::{PyInt, PyIntRef, PyStrRef, PyTypeRef}, function::{ArgBytesLike, FuncArgs, KwArgs, OptionalArg}, iterator, - protocol::PyIter, + protocol::{PyIter, PyIterReturn}, slots::{ Callable, PyComparisonOp::{Eq, Ge, Gt, Le, Lt, Ne}, @@ -220,7 +220,7 @@ mod _operator { #[pyfunction(name = "countOf")] fn count_of(a: PyIter, b: PyObjectRef, vm: &VirtualMachine) -> PyResult { let mut count: usize = 0; - while let Some(element) = iterator::get_next_object(vm, &a)? { + while let PyIterReturn::Return(element) = a.next(vm)? { if element.is(&b) || vm.bool_eq(&b, &element)? { count += 1; } @@ -244,7 +244,7 @@ mod _operator { #[pyfunction(name = "indexOf")] fn index_of(a: PyIter, b: PyObjectRef, vm: &VirtualMachine) -> PyResult { let mut index: usize = 0; - while let Some(element) = iterator::get_next_object(vm, &a)? { + while let PyIterReturn::Return(element) = a.next(vm)? { if element.is(&b) || vm.bool_eq(&b, &element)? { return Ok(index); } diff --git a/vm/src/stdlib/os.rs b/vm/src/stdlib/os.rs index f2f7f675a..d4c4832f8 100644 --- a/vm/src/stdlib/os.rs +++ b/vm/src/stdlib/os.rs @@ -430,6 +430,7 @@ pub(super) mod _os { crt_fd::{Fd, Offset}, exceptions::IntoPyException, function::{ArgBytesLike, FuncArgs, OptionalArg}, + protocol::PyIterReturn, slots::{IteratorIterable, SlotIterator}, suppress_iph, utils::Either, @@ -899,28 +900,30 @@ pub(super) mod _os { } impl IteratorIterable for ScandirIterator {} impl SlotIterator for ScandirIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { if zelf.exhausted.load() { - return Err(vm.new_stop_iteration()); + return Ok(PyIterReturn::StopIteration(None)); } match zelf.entries.write().next() { Some(entry) => match entry { - Ok(entry) => Ok(DirEntry { - entry, - mode: zelf.mode, - lstat: OnceCell::new(), - stat: OnceCell::new(), - #[cfg(not(unix))] - ino: AtomicCell::new(None), - } - .into_ref(vm) - .into_object()), + Ok(entry) => Ok(PyIterReturn::Return( + DirEntry { + entry, + mode: zelf.mode, + lstat: OnceCell::new(), + stat: OnceCell::new(), + #[cfg(not(unix))] + ino: AtomicCell::new(None), + } + .into_ref(vm) + .into_object(), + )), Err(err) => Err(err.into_pyexception(vm)), }, None => { zelf.exhausted.store(true); - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } diff --git a/vm/src/stdlib/pystruct.rs b/vm/src/stdlib/pystruct.rs index 085f73f92..04747d8cf 100644 --- a/vm/src/stdlib/pystruct.rs +++ b/vm/src/stdlib/pystruct.rs @@ -18,6 +18,7 @@ pub(crate) mod _struct { }, common::str::wchar_t, function::{ArgBytesLike, ArgMemoryBuffer, PosArgs}, + protocol::PyIterReturn, slots::{IteratorIterable, SlotConstructor, SlotIterator}, utils::Either, IntoPyObject, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, VirtualMachine, @@ -832,14 +833,16 @@ pub(crate) mod _struct { } impl IteratorIterable for UnpackIterator {} impl SlotIterator for UnpackIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let size = zelf.format_spec.size; let offset = zelf.offset.fetch_add(size); zelf.buffer.with_ref(|buf| { if let Some(buf) = buf.get(offset..offset + size) { - zelf.format_spec.unpack(buf, vm).map(|x| x.into_object()) + zelf.format_spec + .unpack(buf, vm) + .map(|x| PyIterReturn::Return(x.into_object())) } else { - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } }) } diff --git a/vm/src/vm.rs b/vm/src/vm.rs index 5e2a00614..ec1ab4a83 100644 --- a/vm/src/vm.rs +++ b/vm/src/vm.rs @@ -22,6 +22,7 @@ use crate::{ frozen, function::{FuncArgs, IntoFuncArgs}, import, iterator, + protocol::PyIterReturn, scope::Scope, signal::NSIG, slots::PyComparisonOp, @@ -1927,7 +1928,7 @@ impl VirtualMachine { fn _membership_iter_search(&self, haystack: PyObjectRef, needle: PyObjectRef) -> PyResult { let iter = haystack.get_iter(self)?; loop { - if let Some(element) = iterator::get_next_object(self, &iter)? { + if let PyIterReturn::Return(element) = iter.next(self)? { if self.bool_eq(&needle, &element)? { return Ok(self.ctx.new_bool(true)); } else { diff --git a/wasm/lib/src/js_module.rs b/wasm/lib/src/js_module.rs index b88042241..5a5014bc1 100644 --- a/wasm/lib/src/js_module.rs +++ b/wasm/lib/src/js_module.rs @@ -6,15 +6,15 @@ use std::{cell, fmt, future}; use wasm_bindgen::{closure::Closure, prelude::*, JsCast}; use wasm_bindgen_futures::{future_to_promise, JsFuture}; -use rustpython_vm::builtins::PyBaseExceptionRef; -use rustpython_vm::builtins::{PyFloatRef, PyStrRef, PyTypeRef}; -use rustpython_vm::function::{OptionalArg, OptionalOption, PosArgs}; -use rustpython_vm::slots::{IteratorIterable, SlotIterator}; -use rustpython_vm::types::create_simple_type; -use rustpython_vm::VirtualMachine; use rustpython_vm::{ - function::ArgCallable, IntoPyObject, PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue, - TryFromObject, + builtins::{PyBaseExceptionRef, PyFloatRef, PyStrRef, PyTypeRef}, + function::ArgCallable, + function::{OptionalArg, OptionalOption, PosArgs}, + protocol::PyIterReturn, + slots::{IteratorIterable, SlotIterator}, + types::create_simple_type, + IntoPyObject, PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, + VirtualMachine, }; #[wasm_bindgen(inline_js = " @@ -589,8 +589,8 @@ impl AwaitPromise { impl IteratorIterable for AwaitPromise {} impl SlotIterator for AwaitPromise { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - zelf.send(None, vm) + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + PyIterReturn::from_result(zelf.send(None, vm), vm) } }