Use raise_if_stop! macro where possible

ShaharNaveh · ShaharNaveh · commit bf0310282b7d · 2025-07-10T19:12:36.000+03:00
diff --git 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<PyEnumerate> {
 }
 
 impl SelfIter for PyEnumerate {}
+
 impl IterNext for PyEnumerate {
     fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
-        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
@@ -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<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
         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
@@ -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<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
         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);
         }
 
