youknowone · Oct 4, 2021 · Oct 1, 2021 · Oct 1, 2021 · Oct 2, 2021 · Oct 2, 2021
diff --git a/stdlib/src/json.rs b/stdlib/src/json.rs
@@ -5,13 +5,12 @@ mod machinery;
 mod _json {
     use super::machinery;
     use crate::vm::{
-        builtins::PyBaseExceptionRef,
-        builtins::{PyStrRef, PyTypeRef},
+        builtins::{PyBaseExceptionRef, PyStrRef, PyTypeRef},
         function::{FuncArgs, OptionalArg},
-        iterator,
+        protocol::PyIterReturn,
         slots::{Callable, SlotConstructor},
-        IdProtocol, IntoPyObject, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
-        VirtualMachine,
+        IdProtocol, IntoPyObject, IntoPyResult, PyObjectRef, PyRef, PyResult, PyValue,
+        TryFromObject, VirtualMachine,
     };
     use num_bigint::BigInt;
     use std::str::FromStr;
@@ -74,42 +73,48 @@ mod _json {
             idx: usize,
             scan_once: PyObjectRef,
             vm: &VirtualMachine,
-        ) -> PyResult {
-            let c = s
-                .chars()
-                .next()
-                .ok_or_else(|| iterator::stop_iter_with_value(vm.ctx.new_int(idx), vm))?;
+        ) -> PyResult<PyIterReturn> {
+            let c = match s.chars().next() {
+                Some(c) => c,
+                None => return Ok(PyIterReturn::StopIteration(Some(vm.ctx.new_int(idx)))),
+            };
             let next_idx = idx + c.len_utf8();
             match c {
                 '"' => {
                     return scanstring(pystr, next_idx, OptionalArg::Present(self.strict), vm)
-                        .map(|x| x.into_pyobject(vm))
+                        .map(|x| PyIterReturn::Return(x.into_pyobject(vm)))
                 }
                 '{' => {
                     // TODO: parse the object in rust
                     let parse_obj = vm.get_attribute(self.ctx.clone(), "parse_object")?;
-                    return vm.invoke(
-                        &parse_obj,
-                        (
-                            vm.ctx
-                                .new_tuple(vec![pystr.into_object(), vm.ctx.new_int(next_idx)]),
-                            self.strict,
-                            scan_once,
-                            self.object_hook.clone(),
-                            self.object_pairs_hook.clone(),
+                    return PyIterReturn::from_pyresult(
+                        vm.invoke(
+                            &parse_obj,
+                            (
+                                vm.ctx
+                                    .new_tuple(vec![pystr.into_object(), vm.ctx.new_int(next_idx)]),
+                                self.strict,
+                                scan_once,
+                                self.object_hook.clone(),
+                                self.object_pairs_hook.clone(),
+                            ),
                         ),
+                        vm,
                     );
                 }
                 '[' => {
                     // TODO: parse the array in rust
                     let parse_array = vm.get_attribute(self.ctx.clone(), "parse_array")?;
-                    return vm.invoke(
-                        &parse_array,
-                        vec![
-                            vm.ctx
-                                .new_tuple(vec![pystr.into_object(), vm.ctx.new_int(next_idx)]),
-                            scan_once,
-                        ],
+                    return PyIterReturn::from_pyresult(
+                        vm.invoke(
+                            &parse_array,
+                            vec![
+                                vm.ctx
+                                    .new_tuple(vec![pystr.into_object(), vm.ctx.new_int(next_idx)]),
+                                scan_once,
+                            ],
+                        ),
+                        vm,
                     );
                 }
                 _ => {}
@@ -118,7 +123,9 @@ mod _json {
             macro_rules! parse_const {
                 ($s:literal, $val:expr) => {
                     if s.starts_with($s) {
-                        return Ok(vm.ctx.new_tuple(vec![$val, vm.ctx.new_int(idx + $s.len())]));
+                        return Ok(PyIterReturn::Return(
+                            vm.ctx.new_tuple(vec![$val, vm.ctx.new_int(idx + $s.len())]),
+                        ));
                     }
                 };
             }
@@ -128,16 +135,18 @@ mod _json {
             parse_const!("false", vm.ctx.new_bool(false));
 
             if let Some((res, len)) = self.parse_number(s, vm) {
-                return Ok(vm.ctx.new_tuple(vec![res?, vm.ctx.new_int(idx + len)]));
+                return Ok(PyIterReturn::Return(
+                    vm.ctx.new_tuple(vec![res?, vm.ctx.new_int(idx + len)]),
+                ));
             }
 
             macro_rules! parse_constant {
                 ($s:literal) => {
                     if s.starts_with($s) {
-                        return Ok(vm.ctx.new_tuple(vec![
+                        return Ok(PyIterReturn::Return(vm.ctx.new_tuple(vec![
                             vm.invoke(&self.parse_constant, ($s.to_owned(),))?,
                             vm.ctx.new_int(idx + $s.len()),
-                        ]));
+                        ])));
                     }
                 };
             }
@@ -146,7 +155,7 @@ mod _json {
             parse_constant!("Infinity");
             parse_constant!("-Infinity");
 
-            Err(iterator::stop_iter_with_value(vm.ctx.new_int(idx), vm))
+            Ok(PyIterReturn::StopIteration(Some(vm.ctx.new_int(idx))))
         }
 
         fn parse_number(&self, s: &str, vm: &VirtualMachine) -> Option<(PyResult, usize)> {
@@ -185,16 +194,19 @@ mod _json {
             Some((ret, buf.len()))
         }
 
-        fn call(zelf: &PyRef<Self>, pystr: PyStrRef, idx: isize, vm: &VirtualMachine) -> PyResult {
+        fn call(
+            zelf: &PyRef<Self>,
+            pystr: PyStrRef,
+            idx: isize,
+            vm: &VirtualMachine,
+        ) -> PyResult<PyIterReturn> {
             if idx < 0 {
                 return Err(vm.new_value_error("idx cannot be negative".to_owned()));
             }
             let idx = idx as usize;
             let mut chars = pystr.as_str().chars();
-            if idx > 0 {
-                chars
-                    .nth(idx - 1)
-                    .ok_or_else(|| iterator::stop_iter_with_value(vm.ctx.new_int(idx), vm))?;
+            if idx > 0 && chars.nth(idx - 1).is_none() {
+                return Ok(PyIterReturn::StopIteration(Some(vm.ctx.new_int(idx))));
             }
             zelf.parse(
                 chars.as_str(),
@@ -209,7 +221,7 @@ mod _json {
     impl Callable for JsonScanner {
         fn call(zelf: &PyRef<Self>, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
             let (pystr, idx) = args.bind::<(PyStrRef, isize)>(vm)?;
-            JsonScanner::call(zelf, pystr, idx, vm)
+            JsonScanner::call(zelf, pystr, idx, vm).into_pyresult(vm)
         }
     }
 

diff --git a/vm/src/anystr.rs b/vm/src/anystr.rs
@@ -1,8 +1,10 @@
-use crate::builtins::int::PyIntRef;
-use crate::cformat::CFormatString;
-use crate::function::{single_or_tuple_any, OptionalOption, PyIterator};
-use crate::vm::VirtualMachine;
-use crate::{PyObjectRef, PyResult, TryFromObject, TypeProtocol};
+use crate::{
+    builtins::PyIntRef,
+    cformat::CFormatString,
+    function::{single_or_tuple_any, OptionalOption},
+    protocol::PyIterIter,
+    PyObjectRef, PyResult, TryFromObject, TypeProtocol, VirtualMachine,
+};
 use num_traits::{cast::ToPrimitive, sign::Signed};
 use std::str::FromStr;
 
@@ -291,7 +293,7 @@ pub trait AnyStr<'s>: 's {
 
     fn py_join<'a>(
         &self,
-        mut iter: PyIterator<'a, impl AnyStrWrapper<'s, Str = Self> + TryFromObject>,
+        mut iter: PyIterIter<'a, impl AnyStrWrapper<'s, Str = Self> + TryFromObject>,
     ) -> PyResult<Self::Container> {
         let mut joined = if let Some(elem) = iter.next() {
             elem?.as_ref().to_container()

diff --git a/vm/src/builtins/asyncgenerator.rs b/vm/src/builtins/asyncgenerator.rs
@@ -1,7 +1,7 @@
 use super::{PyCode, PyStrRef, PyTypeRef};
 use crate::{
     builtins::PyBaseExceptionRef,
-    coroutine::{Coro, Variant},
+    coroutine::Coro,
     frame::FrameRef,
     function::OptionalArg,
     protocol::PyIterReturn,
@@ -34,7 +34,7 @@ impl PyAsyncGen {
 
     pub fn new(frame: FrameRef, name: PyStrRef) -> Self {
         PyAsyncGen {
-            inner: Coro::new(frame, Variant::AsyncGen, name),
+            inner: Coro::new(frame, name),
             running_async: AtomicCell::new(false),
         }
     }
@@ -50,8 +50,8 @@ impl PyAsyncGen {
     }
 
     #[pymethod(magic)]
-    fn repr(zelf: PyRef<Self>) -> String {
-        zelf.inner.repr(zelf.get_id())
+    fn repr(zelf: PyRef<Self>, vm: &VirtualMachine) -> String {
+        zelf.inner.repr(zelf.as_object(), zelf.get_id(), vm)
     }
 
     #[pymethod(magic)]
@@ -138,24 +138,24 @@ impl PyValue for PyAsyncGenWrappedValue {
 impl PyAsyncGenWrappedValue {}
 
 impl PyAsyncGenWrappedValue {
-    fn unbox(ag: &PyAsyncGen, val: PyResult, vm: &VirtualMachine) -> PyResult {
-        if let Err(ref e) = val {
-            if e.isinstance(&vm.ctx.exceptions.stop_async_iteration)
-                || e.isinstance(&vm.ctx.exceptions.generator_exit)
-            {
-                ag.inner.closed.store(true);
-            }
+    fn unbox(ag: &PyAsyncGen, val: PyResult<PyIterReturn>, vm: &VirtualMachine) -> PyResult {
+        let (closed, async_done) = match &val {
+            Ok(PyIterReturn::StopIteration(_)) => (true, true),
+            Err(e) if e.isinstance(&vm.ctx.exceptions.generator_exit) => (true, true),
+            Err(_) => (false, true),
+            _ => (false, false),
+        };
+        if closed {
+            ag.inner.closed.store(true);
+        }
+        if async_done {
             ag.running_async.store(false);
         }
-        let val = val?;
-
+        let val = val?.into_async_pyresult(vm)?;
         match_class!(match val {
             val @ Self => {
                 ag.running_async.store(false);
-                Err(vm.new_exception(
-                    vm.ctx.exceptions.stop_iteration.clone(),
-                    vec![val.0.clone()],
-                ))
+                Err(vm.new_stop_iteration(Some(val.0.clone())))
             }
             val => Ok(val),
         })
@@ -214,7 +214,7 @@ impl PyAsyncGenASend {
                 }
             }
         };
-        let res = self.ag.inner.send(val, vm);
+        let res = self.ag.inner.send(self.ag.as_object(), val, vm);
         let res = PyAsyncGenWrappedValue::unbox(&self.ag, res, vm);
         if res.is_err() {
             self.close();
@@ -237,6 +237,7 @@ impl PyAsyncGenASend {
         }
 
         let res = self.ag.inner.throw(
+            self.ag.as_object(),
             exc_type,
             exc_val.unwrap_or_none(vm),
             exc_tb.unwrap_or_none(vm),
@@ -258,8 +259,7 @@ impl PyAsyncGenASend {
 impl IteratorIterable for PyAsyncGenASend {}
 impl SlotIterator for PyAsyncGenASend {
     fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
-        // TODO: Fix zelf.send to return PyIterReturn
-        PyIterReturn::from_result(zelf.send(vm.ctx.none(), vm), vm)
+        PyIterReturn::from_pyresult(zelf.send(vm.ctx.none(), vm), vm)
     }
 }
 
@@ -304,7 +304,7 @@ impl PyAsyncGenAThrow {
                 }
                 if self.ag.inner.closed() {
                     self.state.store(AwaitableState::Closed);
-                    return Err(vm.new_exception_empty(vm.ctx.exceptions.stop_iteration.clone()));
+                    return Err(vm.new_stop_iteration(None));
                 }
                 if !vm.is_none(&val) {
                     return Err(vm.new_runtime_error(
@@ -315,27 +315,28 @@ impl PyAsyncGenAThrow {
                 self.ag.running_async.store(true);
 
                 let (ty, val, tb) = self.value.clone();
-                let ret = self.ag.inner.throw(ty, val, tb, vm);
+                let ret = self.ag.inner.throw(self.ag.as_object(), ty, val, tb, vm);
                 let ret = if self.aclose {
                     if self.ignored_close(&ret) {
                         Err(self.yield_close(vm))
                     } else {
-                        ret
+                        ret.and_then(|o| o.into_async_pyresult(vm))
                     }
                 } else {
                     PyAsyncGenWrappedValue::unbox(&self.ag, ret, vm)
                 };
                 ret.map_err(|e| self.check_error(e, vm))
             }
             AwaitableState::Iter => {
-                let ret = self.ag.inner.send(val, vm);
+                let ret = self.ag.inner.send(self.ag.as_object(), val, vm);
                 if self.aclose {
                     match ret {
-                        Ok(v) if v.payload_is::<PyAsyncGenWrappedValue>() => {
+                        Ok(PyIterReturn::Return(v)) if v.payload_is::<PyAsyncGenWrappedValue>() => {
                             Err(self.yield_close(vm))
                         }
-                        Ok(v) => Ok(v),
-                        Err(e) => Err(self.check_error(e, vm)),
+                        other => other
+                            .and_then(|o| o.into_async_pyresult(vm))
+                            .map_err(|e| self.check_error(e, vm)),
                     }
                 } else {
                     PyAsyncGenWrappedValue::unbox(&self.ag, ret, vm)
@@ -353,6 +354,7 @@ impl PyAsyncGenAThrow {
         vm: &VirtualMachine,
     ) -> PyResult {
         let ret = self.ag.inner.throw(
+            self.ag.as_object(),
             exc_type,
             exc_val.unwrap_or_none(vm),
             exc_tb.unwrap_or_none(vm),
@@ -362,7 +364,7 @@ impl PyAsyncGenAThrow {
             if self.ignored_close(&ret) {
                 Err(self.yield_close(vm))
             } else {
-                ret
+                ret.and_then(|o| o.into_async_pyresult(vm))
             }
         } else {
             PyAsyncGenWrappedValue::unbox(&self.ag, ret, vm)
@@ -375,9 +377,11 @@ impl PyAsyncGenAThrow {
         self.state.store(AwaitableState::Closed);
     }
 
-    fn ignored_close(&self, res: &PyResult) -> bool {
-        res.as_ref()
-            .map_or(false, |v| v.payload_is::<PyAsyncGenWrappedValue>())
+    fn ignored_close(&self, res: &PyResult<PyIterReturn>) -> bool {
+        res.as_ref().map_or(false, |v| match v {
+            PyIterReturn::Return(obj) => obj.payload_is::<PyAsyncGenWrappedValue>(),
+            PyIterReturn::StopIteration(_) => false,
+        })
     }
     fn yield_close(&self, vm: &VirtualMachine) -> PyBaseExceptionRef {
         self.ag.running_async.store(false);
@@ -391,7 +395,7 @@ impl PyAsyncGenAThrow {
             && (exc.isinstance(&vm.ctx.exceptions.stop_async_iteration)
                 || exc.isinstance(&vm.ctx.exceptions.generator_exit))
         {
-            vm.new_exception_empty(vm.ctx.exceptions.stop_iteration.clone())
+            vm.new_stop_iteration(None)
         } else {
             exc
         }
@@ -401,8 +405,7 @@ impl PyAsyncGenAThrow {
 impl IteratorIterable for PyAsyncGenAThrow {}
 impl SlotIterator for PyAsyncGenAThrow {
     fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
-        // TODO: Fix zelf.send to return PyIterReturn
-        PyIterReturn::from_result(zelf.send(vm.ctx.none(), vm), vm)
+        PyIterReturn::from_pyresult(zelf.send(vm.ctx.none(), vm), vm)
     }
 }