From c31b3b2af940d086dff4e0e3127331e2ccd3a73e Mon Sep 17 00:00:00 2001
From: Noah <33094578+coolreader18@users.noreply.github.com>
Date: Mon, 14 Oct 2019 21:05:39 +0000
Subject: [PATCH] Use the first arg of StopIteration as `yield from` return
 value

---
 tests/snippets/generators.py | 16 ++++++++++
 vm/src/frame.rs              | 62 ++++++++++++++++++++++++++++++++----
 vm/src/obj/objgenerator.rs   | 19 ++---------
 3 files changed, 75 insertions(+), 22 deletions(-)

diff --git a/tests/snippets/generators.py b/tests/snippets/generators.py
index f23e5918a8..463bfc313d 100644
--- a/tests/snippets/generators.py
+++ b/tests/snippets/generators.py
@@ -88,3 +88,19 @@ l = list(g)
 # print(l)
 assert l == [99]
 assert r == ['a', 66, None]
+
+def binary(n):
+    if n <= 1:
+        return 1
+    l = yield from binary(n - 1)
+    r = yield from binary(n - 1)
+    return l + 1 + r
+
+with assert_raises(StopIteration):
+    try:
+        next(binary(5))
+    except StopIteration as stopiter:
+        # TODO: StopIteration.value
+        assert stopiter.args[0] == 31
+        raise
+
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
index b5e3eca169..46e1698e9c 100644
--- a/vm/src/frame.rs
+++ b/vm/src/frame.rs
@@ -105,6 +105,45 @@ pub enum ExecutionResult {
     Yield(PyObjectRef),
 }
 
+impl ExecutionResult {
+    /// Extract an ExecutionResult from a PyResult returned from e.g. gen.__next__() or gen.send()
+    pub fn from_result(vm: &VirtualMachine, res: PyResult) -> PyResult<Self> {
+        match res {
+            Ok(val) => Ok(ExecutionResult::Yield(val)),
+            Err(err) => {
+                if objtype::isinstance(&err, &vm.ctx.exceptions.stop_iteration) {
+                    let args = vm.get_attribute(err, "args")?;
+                    let args: &PyTuple = args.payload().unwrap();
+                    let val = if args.elements.len() == 0 {
+                        vm.get_none()
+                    } else {
+                        args.elements[0].clone()
+                    };
+                    Ok(ExecutionResult::Return(val))
+                } else {
+                    Err(err)
+                }
+            }
+        }
+    }
+
+    /// Turn an ExecutionResult into a PyResult that would be returned from a generator or coroutine
+    pub fn into_result(self, vm: &VirtualMachine) -> PyResult {
+        match self {
+            ExecutionResult::Yield(value) => Ok(value),
+            ExecutionResult::Return(value) => {
+                let stop_iteration = vm.ctx.exceptions.stop_iteration.clone();
+                let args = if vm.is_none(&value) {
+                    vec![]
+                } else {
+                    vec![value]
+                };
+                Err(vm.new_exception_obj(stop_iteration, args).unwrap())
+            }
+        }
+    }
+}
+
 /// A valid execution result, or an exception
 pub type FrameResult = PyResult<Option<ExecutionResult>>;
 
@@ -956,18 +995,29 @@ impl Frame {
 
     fn execute_yield_from(&self, vm: &VirtualMachine) -> FrameResult {
         // Value send into iterator:
-        self.pop_value();
+        let val = self.pop_value();
 
-        let top_of_stack = self.last_value();
-        let next_obj = objiter::get_next_object(vm, &top_of_stack)?;
+        let coro = self.last_value();
 
-        match next_obj {
-            Some(value) => {
+        let result = if vm.is_none(&val) {
+            objiter::call_next(vm, &coro)
+        } else {
+            vm.call_method(&coro, "send", vec![val])
+        };
+
+        let result = ExecutionResult::from_result(vm, result)?;
+
+        match result {
+            ExecutionResult::Yield(value) => {
                 // Set back program counter:
                 *self.lasti.borrow_mut() -= 1;
                 Ok(Some(ExecutionResult::Yield(value)))
             }
-            None => Ok(None),
+            ExecutionResult::Return(value) => {
+                self.pop_value();
+                self.push_value(value);
+                Ok(None)
+            }
         }
     }
 
diff --git a/vm/src/obj/objgenerator.rs b/vm/src/obj/objgenerator.rs
index 8da98c4ef1..4518d164af 100644
--- a/vm/src/obj/objgenerator.rs
+++ b/vm/src/obj/objgenerator.rs
@@ -3,7 +3,7 @@
  */
 
 use super::objtype::{isinstance, PyClassRef};
-use crate::frame::{ExecutionResult, FrameRef};
+use crate::frame::FrameRef;
 use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
 use crate::vm::VirtualMachine;
 
@@ -41,8 +41,7 @@ impl PyGenerator {
     fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult {
         self.frame.push_value(value.clone());
 
-        let result = vm.run_frame(self.frame.clone())?;
-        handle_execution_result(result, vm)
+        vm.run_frame(self.frame.clone())?.into_result(vm)
     }
 
     #[pymethod]
@@ -58,19 +57,7 @@ impl PyGenerator {
         if !isinstance(&exc_val, &vm.ctx.exceptions.base_exception_type) {
             return Err(vm.new_type_error("Can't throw non exception".to_string()));
         }
-        let result = vm.frame_throw(self.frame.clone(), exc_val)?;
-        handle_execution_result(result, vm)
-    }
-}
-
-fn handle_execution_result(result: ExecutionResult, vm: &VirtualMachine) -> PyResult {
-    match result {
-        ExecutionResult::Yield(value) => Ok(value),
-        ExecutionResult::Return(_value) => {
-            // Stop iteration!
-            let stop_iteration = vm.ctx.exceptions.stop_iteration.clone();
-            Err(vm.new_exception(stop_iteration, "End of generator".to_string()))
-        }
+        vm.frame_throw(self.frame.clone(), exc_val)?.into_result(vm)
     }
 }