rebase to use PyIterReturn

vm/src/builtins/bytearray.rs

@@ -751,7 +751,7 @@ impl PyByteArrayIterator {
     fn reduce(&self, vm: &VirtualMachine) -> PyObjectRef {
         self.internal
             .lock()
-            .builtin_iter_reduce(|x| x.clone().into_object(), vm)
+            .builtins_iter_reduce(|x| x.clone().into_object(), vm)
     }
 
     #[pymethod(magic)]
@@ -763,15 +763,13 @@ impl PyByteArrayIterator {
 }
 impl IteratorIterable for PyByteArrayIterator {}
 impl SlotIterator for PyByteArrayIterator {
-    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
-        zelf.internal.lock().next(
-            |bytearray, pos| {
-                let buf = bytearray.borrow_buf();
-                buf.get(pos)
-                    .ok_or_else(|| vm.new_stop_iteration())
-                    .map(|&x| vm.ctx.new_int(x))
-            },
-            vm,
-        )
+    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
+        zelf.internal.lock().next(|bytearray, pos| {
+            let buf = bytearray.borrow_buf();
+            Ok(match buf.get(pos) {
+                Some(&x) => PyIterReturn::Return(vm.ctx.new_int(x)),
+                None => PyIterReturn::StopIteration(None),
+            })
+        })
     }
 }

vm/src/builtins/bytes.rs

@@ -605,7 +605,7 @@ impl PyBytesIterator {
     fn reduce(&self, vm: &VirtualMachine) -> PyObjectRef {
         self.internal
             .lock()
-            .builtin_iter_reduce(|x| x.clone().into_object(), vm)
+            .builtins_iter_reduce(|x| x.clone().into_object(), vm)
     }
 
     #[pymethod(magic)]
@@ -617,17 +617,13 @@ impl PyBytesIterator {
 }
 impl IteratorIterable for PyBytesIterator {}
 impl SlotIterator for PyBytesIterator {
-    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
-        zelf.internal.lock().next(
-            |bytes, pos| {
-                bytes
-                    .as_bytes()
-                    .get(pos)
-                    .ok_or_else(|| vm.new_stop_iteration())
-                    .map(|&x| vm.ctx.new_int(x))
-            },
-            vm,
-        )
+    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
+        zelf.internal.lock().next(|bytes, pos| {
+            Ok(match bytes.as_bytes().get(pos) {
+                Some(&x) => PyIterReturn::Return(vm.ctx.new_int(x)),
+                None => PyIterReturn::StopIteration(None),
+            })
+        })
     }
 }
 

vm/src/builtins/dict.rs

@@ -733,7 +733,7 @@ macro_rules! dict_iterator {
         impl IteratorIterable for $iter_name {}
         impl SlotIterator for $iter_name {
             #[allow(clippy::redundant_closure_call)]
-            fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
+            fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
                 let mut internal = zelf.internal.lock();
                 if let IterStatus::Active(dict) = &internal.status {
                     if dict.entries.has_changed_size(&zelf.size) {
@@ -745,15 +745,15 @@ macro_rules! dict_iterator {
                     match dict.entries.next_entry(internal.position) {
                         Some((position, key, value)) => {
                             internal.position = position;
-                            Ok(($result_fn)(vm, key, value))
+                            Ok(PyIterReturn::Return(($result_fn)(vm, key, value)))
                         }
                         None => {
                             internal.status = IterStatus::Exhausted;
-                            Err(vm.new_stop_iteration())
+                            Ok(PyIterReturn::StopIteration(None))
                         }
                     }
                 } else {
-                    Err(vm.new_stop_iteration())
+                    Ok(PyIterReturn::StopIteration(None))
                 }
             }
         }
@@ -793,7 +793,7 @@ 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<Self>, vm: &VirtualMachine) -> PyResult {
+            fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
                 let mut internal = zelf.internal.lock();
                 if let IterStatus::Active(dict) = &internal.status {
                     if dict.entries.has_changed_size(&zelf.size) {
@@ -809,15 +809,15 @@ macro_rules! dict_iterator {
                             } else {
                                 internal.position = position;
                             }
-                            Ok(($result_fn)(vm, key, value))
+                            Ok(PyIterReturn::Return(($result_fn)(vm, key, value)))
                         }
                         None => {
                             internal.status = IterStatus::Exhausted;
-                            Err(vm.new_stop_iteration())
+                            Ok(PyIterReturn::StopIteration(None))
                         }
                     }
                 } else {
-                    Err(vm.new_stop_iteration())
+                    Ok(PyIterReturn::StopIteration(None))
                 }
             }
         }

vm/src/builtins/enumerate.rs

@@ -1,5 +1,6 @@
 use super::{IterStatus, PositionIterInternal, PyIntRef, PyTypeRef};
 use crate::common::lock::{PyMutex, PyRwLock};
+use crate::TypeProtocol;
 use crate::{
     function::OptionalArg,
     protocol::{PyIter, PyIterReturn},
@@ -105,16 +106,26 @@ impl PyReverseSequenceIterator {
     fn reduce(&self, vm: &VirtualMachine) -> PyObjectRef {
         self.internal
             .lock()
-            .builtin_reversed_reduce(|x| x.clone(), vm)
+            .builtins_reversed_reduce(|x| x.clone(), vm)
     }
 }
 
 impl IteratorIterable for PyReverseSequenceIterator {}
 impl SlotIterator for PyReverseSequenceIterator {
-    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
+    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
         zelf.internal
             .lock()
-            .rev_next(|obj, pos| obj.get_item(pos, vm), vm)
+            .rev_next(|obj, pos| match obj.get_item(pos, vm) {
+                Ok(ret) => Ok(PyIterReturn::Return(ret)),
+                Err(e) if e.isinstance(&vm.ctx.exceptions.index_error) => {
+                    Ok(PyIterReturn::StopIteration(None))
+                }
+                Err(e) if e.isinstance(&vm.ctx.exceptions.stop_iteration) => {
+                    let args = e.get_arg(0);
+                    Ok(PyIterReturn::StopIteration(args))
+                }
+                Err(e) => Err(e),
+            })
     }
 }
 

vm/src/builtins/iter.rs

@@ -71,7 +71,7 @@ impl<T> PositionIterInternal<T> {
         }
     }
 
-    pub fn builtin_iter_reduce<F>(&self, f: F, vm: &VirtualMachine) -> PyObjectRef
+    pub fn builtins_iter_reduce<F>(&self, f: F, vm: &VirtualMachine) -> PyObjectRef
     where
         F: FnOnce(&T) -> PyObjectRef,
     {
@@ -79,7 +79,7 @@ impl<T> PositionIterInternal<T> {
         self._reduce(iter, f, vm)
     }
 
-    pub fn builtin_reversed_reduce<F>(&self, f: F, vm: &VirtualMachine) -> PyObjectRef
+    pub fn builtins_reversed_reduce<F>(&self, f: F, vm: &VirtualMachine) -> PyObjectRef
     where
         F: FnOnce(&T) -> PyObjectRef,
     {
@@ -87,66 +87,42 @@ impl<T> PositionIterInternal<T> {
         self._reduce(reversed, f, vm)
     }
 
-    pub fn next<F>(&mut self, f: F, vm: &VirtualMachine) -> PyResult
+    fn _next<F, OP>(&mut self, f: F, op: OP) -> PyResult<PyIterReturn>
     where
-        F: FnOnce(&T, usize) -> PyResult,
+        F: FnOnce(&T, usize) -> PyResult<PyIterReturn>,
+        OP: FnOnce(&mut Self),
     {
         if let IterStatus::Active(obj) = &self.status {
-            match f(obj, self.position) {
-                Err(e) if e.isinstance(&vm.ctx.exceptions.stop_iteration) => {
-                    self.status = IterStatus::Exhausted;
-                    Err(e)
-                }
-                Err(e) if e.isinstance(&vm.ctx.exceptions.index_error) => {
-                    self.status = IterStatus::Exhausted;
-                    Err(vm.new_stop_iteration())
-                }
-                Err(e) if e.isinstance(&vm.ctx.exceptions.runtime_error) => {
-                    self.status = IterStatus::Exhausted;
-                    Err(e)
-                }
-                Err(e) => Err(e),
-                Ok(ret) => {
-                    self.position += 1;
-                    Ok(ret)
-                }
+            let ret = f(obj, self.position);
+            if let Ok(PyIterReturn::Return(_)) = ret {
+                op(self);
+            } else {
+                self.status = IterStatus::Exhausted;
             }
+            ret
         } else {
-            Err(vm.new_stop_iteration())
+            Ok(PyIterReturn::StopIteration(None))
         }
     }
 
-    pub fn rev_next<F>(&mut self, f: F, vm: &VirtualMachine) -> PyResult
+    pub fn next<F>(&mut self, f: F) -> PyResult<PyIterReturn>
     where
-        F: FnOnce(&T, usize) -> PyResult,
+        F: FnOnce(&T, usize) -> PyResult<PyIterReturn>,
     {
-        if let IterStatus::Active(obj) = &self.status {
-            match f(obj, self.position) {
-                Err(e) if e.isinstance(&vm.ctx.exceptions.stop_iteration) => {
-                    self.status = IterStatus::Exhausted;
-                    Err(e)
-                }
-                Err(e) if e.isinstance(&vm.ctx.exceptions.index_error) => {
-                    self.status = IterStatus::Exhausted;
-                    Err(vm.new_stop_iteration())
-                }
-                Err(e) if e.isinstance(&vm.ctx.exceptions.runtime_error) => {
-                    self.status = IterStatus::Exhausted;
-                    Err(e)
-                }
-                Err(e) => Err(e),
-                Ok(ret) => {
-                    if self.position == 0 {
-                        self.status = IterStatus::Exhausted;
-                    } else {
-                        self.position -= 1;
-                    }
-                    Ok(ret)
-                }
+        self._next(f, |zelf| zelf.position += 1)
+    }
+
+    pub fn rev_next<F>(&mut self, f: F) -> PyResult<PyIterReturn>
+    where
+        F: FnOnce(&T, usize) -> PyResult<PyIterReturn>,
+    {
+        self._next(f, |zelf| {
+            if zelf.position == 0 {
+                zelf.status = IterStatus::Exhausted;
+            } else {
+                zelf.position -= 1;
             }
-        } else {
-            Err(vm.new_stop_iteration())
-        }
+        })
     }
 
     pub fn length_hint<F>(&self, f: F) -> usize
@@ -221,7 +197,7 @@ impl PySequenceIterator {
 
     #[pymethod(magic)]
     fn reduce(&self, vm: &VirtualMachine) -> PyObjectRef {
-        self.internal.lock().builtin_iter_reduce(|x| x.clone(), vm)
+        self.internal.lock().builtins_iter_reduce(|x| x.clone(), vm)
     }
 
     #[pymethod(magic)]
@@ -232,10 +208,20 @@ impl PySequenceIterator {
 
 impl IteratorIterable for PySequenceIterator {}
 impl SlotIterator for PySequenceIterator {
-    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
+    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
         zelf.internal
             .lock()
-            .next(|obj, pos| obj.get_item(pos, vm), vm)
+            .next(|obj, pos| match obj.get_item(pos, vm) {
+                Ok(ret) => Ok(PyIterReturn::Return(ret)),
+                Err(e) if e.isinstance(&vm.ctx.exceptions.index_error) => {
+                    Ok(PyIterReturn::StopIteration(None))
+                }
+                Err(e) if e.isinstance(&vm.ctx.exceptions.stop_iteration) => {
+                    let args = e.get_arg(0);
+                    Ok(PyIterReturn::StopIteration(args))
+                }
+                Err(e) => Err(e),
+            })
     }
 }
 
@@ -264,18 +250,18 @@ impl PyCallableIterator {
 
 impl IteratorIterable for PyCallableIterator {}
 impl SlotIterator for PyCallableIterator {
-    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
+    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
         let status = zelf.status.upgradable_read();
         if let IterStatus::Active(callable) = &*status {
             let ret = callable.invoke((), vm)?;
             if vm.bool_eq(&ret, &zelf.sentinel)? {
                 *PyRwLockUpgradableReadGuard::upgrade(status) = IterStatus::Exhausted;
-                Err(vm.new_stop_iteration())
+                Ok(PyIterReturn::StopIteration(None))
             } else {
-                Ok(ret)
+                Ok(PyIterReturn::Return(ret))
             }
         } else {
-            Err(vm.new_stop_iteration())
+            Ok(PyIterReturn::StopIteration(None))
         }
     }
 }

vm/src/builtins/list.rs

@@ -489,22 +489,20 @@ impl PyListIterator {
     fn reduce(&self, vm: &VirtualMachine) -> PyObjectRef {
         self.internal
             .lock()
-            .builtin_iter_reduce(|x| x.clone().into_object(), vm)
+            .builtins_iter_reduce(|x| x.clone().into_object(), vm)
     }
 }
 
 impl IteratorIterable for PyListIterator {}
 impl SlotIterator for PyListIterator {
-    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
-        zelf.internal.lock().next(
-            |list, pos| {
-                let vec = list.borrow_vec();
-                vec.get(pos)
-                    .ok_or_else(|| vm.new_stop_iteration())
-                    .map(|x| x.clone())
-            },
-            vm,
-        )
+    fn next(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<PyIterReturn> {
+        zelf.internal.lock().next(|list, pos| {
+            let vec = list.borrow_vec();
+            Ok(match vec.get(pos) {
+                Some(x) => PyIterReturn::Return(x.clone()),
+                None => PyIterReturn::StopIteration(None),
+            })
+        })
     }
 }
 
@@ -538,22 +536,20 @@ impl PyListReverseIterator {
     fn reduce(&self, vm: &VirtualMachine) -> PyObjectRef {
         self.internal
             .lock()
-            .builtin_reversed_reduce(|x| x.clone().into_object(), vm)
+            .builtins_reversed_reduce(|x| x.clone().into_object(), vm)
     }
 }
 
 impl IteratorIterable for PyListReverseIterator {}
 impl SlotIterator for PyListReverseIterator {
-    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
-        zelf.internal.lock().rev_next(
-            |list, pos| {
-                let vec = list.borrow_vec();
-                vec.get(pos)
-                    .ok_or_else(|| vm.new_stop_iteration())
-                    .map(|x| x.clone())
-            },
-            vm,
-        )
+    fn next(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<PyIterReturn> {
+        zelf.internal.lock().rev_next(|list, pos| {
+            let vec = list.borrow_vec();
+            Ok(match vec.get(pos) {
+                Some(x) => PyIterReturn::Return(x.clone()),
+                None => PyIterReturn::StopIteration(None),
+            })
+        })
     }
 }
 

vm/src/builtins/pystr.rs

@@ -200,13 +200,13 @@ impl PyStrIterator {
         self.internal
             .lock()
             .0
-            .builtin_iter_reduce(|x| x.clone().into_object(), vm)
+            .builtins_iter_reduce(|x| x.clone().into_object(), vm)
     }
 }
 
 impl IteratorIterable for PyStrIterator {}
 impl SlotIterator for PyStrIterator {
-    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
+    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
         let mut internal = zelf.internal.lock();
 
         if let IterStatus::Active(s) = &internal.0.status {
@@ -216,20 +216,18 @@ impl SlotIterator for PyStrIterator {
                 if let Some((offset, ch)) = value.char_indices().nth(internal.0.position) {
                     internal.0.position += 1;
                     internal.1 = offset + ch.len_utf8();
-                    return Ok(ch.into_pyobject(vm));
+                    return Ok(PyIterReturn::Return(ch.into_pyobject(vm)));
                 }
             } else if let Some(value) = value.get(internal.1..) {
                 if let Some(ch) = value.chars().next() {
                     internal.0.position += 1;
                     internal.1 += ch.len_utf8();
-                    return Ok(ch.into_pyobject(vm));
+                    return Ok(PyIterReturn::Return(ch.into_pyobject(vm)));
                 }
             }
             internal.0.status = Exhausted;
-            Err(vm.new_stop_iteration())
-        } else {
-            Err(vm.new_stop_iteration())
         }
+        Ok(PyIterReturn::StopIteration(None))
     }
 }
 

vm/src/builtins/set.rs

@@ -853,7 +853,7 @@ impl PySetIterator {
 
 impl IteratorIterable for PySetIterator {}
 impl SlotIterator for PySetIterator {
-    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
+    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
         let mut internal = zelf.internal.lock();
         if let IterStatus::Active(dict) = &internal.status {
             if dict.has_changed_size(&zelf.size) {
@@ -863,15 +863,15 @@ impl SlotIterator for PySetIterator {
             match dict.next_entry(internal.position) {
                 Some((position, key, _)) => {
                     internal.position = position;
-                    Ok(key)
+                    Ok(PyIterReturn::Return(key))
                 }
                 None => {
                     internal.status = IterStatus::Exhausted;
-                    Err(vm.new_stop_iteration())
+                    Ok(PyIterReturn::StopIteration(None))
                 }
             }
         } else {
-            Err(vm.new_stop_iteration())
+            Ok(PyIterReturn::StopIteration(None))
         }
     }
 }

vm/src/builtins/tuple.rs

@@ -342,23 +342,20 @@ impl PyTupleIterator {
     fn reduce(&self, vm: &VirtualMachine) -> PyObjectRef {
         self.internal
             .lock()
-            .builtin_iter_reduce(|x| x.clone().into_object(), vm)
+            .builtins_iter_reduce(|x| x.clone().into_object(), vm)
     }
 }
 
 impl IteratorIterable for PyTupleIterator {}
 impl SlotIterator for PyTupleIterator {
-    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
-        zelf.internal.lock().next(
-            |tuple, pos| {
-                tuple
-                    .as_slice()
-                    .get(pos)
-                    .ok_or_else(|| vm.new_stop_iteration())
-                    .map(|x| x.clone())
-            },
-            vm,
-        )
+    fn next(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<PyIterReturn> {
+        zelf.internal.lock().next(|tuple, pos| {
+            Ok(if let Some(ret) = tuple.as_slice().get(pos) {
+                PyIterReturn::Return(ret.clone())
+            } else {
+                PyIterReturn::StopIteration(None)
+            })
+        })
     }
 }
 

vm/src/stdlib/collections.rs

@@ -634,22 +634,17 @@ mod _collections {
 
     impl IteratorIterable for PyDequeIterator {}
     impl SlotIterator for PyDequeIterator {
-        fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
-            zelf.internal.lock().next(
-                |deque, pos| {
-                    if zelf.state != deque.state.load() {
-                        return Err(
-                            vm.new_runtime_error("Deque mutated during iteration".to_owned())
-                        );
-                    }
-                    let deque = deque.borrow_deque();
-                    deque
-                        .get(pos)
-                        .ok_or_else(|| vm.new_stop_iteration())
-                        .map(|x| x.clone())
-                },
-                vm,
-            )
+        fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
+            zelf.internal.lock().next(|deque, pos| {
+                if zelf.state != deque.state.load() {
+                    return Err(vm.new_runtime_error("Deque mutated during iteration".to_owned()));
+                }
+                let deque = deque.borrow_deque();
+                Ok(match deque.get(pos) {
+                    Some(x) => PyIterReturn::Return(x.clone()),
+                    None => PyIterReturn::StopIteration(None),
+                })
+            })
         }
     }
 
@@ -706,26 +701,23 @@ mod _collections {
 
     impl IteratorIterable for PyReverseDequeIterator {}
     impl SlotIterator for PyReverseDequeIterator {
-        fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
-            zelf.internal.lock().next(
-                |deque, pos| {
-                    if deque.state.load() != zelf.state {
-                        return Err(
-                            vm.new_runtime_error("Deque mutated during iteration".to_owned())
-                        );
-                    }
-                    let deque = deque.borrow_deque();
-                    let pos = deque
+        fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
+            zelf.internal.lock().next(|deque, pos| {
+                if deque.state.load() != zelf.state {
+                    return Err(vm.new_runtime_error("Deque mutated during iteration".to_owned()));
+                }
+                let deque = deque.borrow_deque();
+                Ok(
+                    match deque
                         .len()
                         .checked_sub(pos + 1)
-                        .ok_or_else(|| vm.new_stop_iteration())?;
-                    deque
-                        .get(pos)
-                        .ok_or_else(|| vm.new_stop_iteration())
-                        .map(|x| x.clone())
-                },
-                vm,
-            )
+                        .and_then(|pos| deque.get(pos))
+                    {
+                        Some(x) => PyIterReturn::Return(x.clone()),
+                        None => PyIterReturn::StopIteration(None),
+                    },
+                )
+            })
         }
     }
 }