diff --git a/vm/src/builtins/enumerate.rs b/vm/src/builtins/enumerate.rs index 64cd4e774..db3d45b24 100644 --- a/vm/src/builtins/enumerate.rs +++ b/vm/src/builtins/enumerate.rs @@ -8,6 +8,7 @@ use crate::{ convert::ToPyObject, function::OptionalArg, protocol::{PyIter, PyIterReturn}, + raise_if_stop, types::{Constructor, IterNext, Iterable, SelfIter}, }; use malachite_bigint::BigInt; @@ -73,12 +74,10 @@ impl Py { } impl SelfIter for PyEnumerate {} + impl IterNext for PyEnumerate { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { - let next_obj = match zelf.iterator.next(vm)? { - PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), - PyIterReturn::Return(obj) => obj, - }; + let next_obj = raise_if_stop!(zelf.iterator.next(vm)?); let mut counter = zelf.counter.write(); let position = counter.clone(); *counter += 1; diff --git a/vm/src/builtins/filter.rs b/vm/src/builtins/filter.rs index 5dd0162f8..661fbd022 100644 --- a/vm/src/builtins/filter.rs +++ b/vm/src/builtins/filter.rs @@ -3,6 +3,7 @@ use crate::{ Context, Py, PyObjectRef, PyPayload, PyResult, VirtualMachine, class::PyClassImpl, protocol::{PyIter, PyIterReturn}, + raise_if_stop, types::{Constructor, IterNext, Iterable, SelfIter}, }; @@ -45,24 +46,22 @@ impl PyFilter { } impl SelfIter for PyFilter {} + impl IterNext for PyFilter { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { let predicate = &zelf.predicate; loop { - let next_obj = match zelf.iterator.next(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), - }; + let next_obj = raise_if_stop!(zelf.iterator.next(vm)?); let predicate_value = if vm.is_none(predicate) { next_obj.clone() } else { - // the predicate itself can raise StopIteration which does stop the filter - // iteration - match PyIterReturn::from_pyresult(predicate.call((next_obj.clone(),), vm), vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), - } + // the predicate itself can raise StopIteration which does stop the filter iteration + raise_if_stop!(PyIterReturn::from_pyresult( + predicate.call((next_obj.clone(),), vm), + vm + )?) }; + if predicate_value.try_to_bool(vm)? { return Ok(PyIterReturn::Return(next_obj)); } diff --git a/vm/src/builtins/map.rs b/vm/src/builtins/map.rs index 004028c2c..06a533f8b 100644 --- a/vm/src/builtins/map.rs +++ b/vm/src/builtins/map.rs @@ -5,6 +5,7 @@ use crate::{ class::PyClassImpl, function::PosArgs, protocol::{PyIter, PyIterReturn}, + raise_if_stop, types::{Constructor, IterNext, Iterable, SelfIter}, }; @@ -53,14 +54,12 @@ impl PyMap { } impl SelfIter for PyMap {} + impl IterNext for PyMap { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { let mut next_objs = Vec::new(); for iterator in &zelf.iterators { - let item = match iterator.next(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), - }; + let item = raise_if_stop!(iterator.next(vm)?); next_objs.push(item); }