Move generator functionality to objcoroinner

vm/src/frame.rs

@@ -1010,14 +1010,15 @@ impl Frame {
     }
 
     fn _send(&self, coro: PyObjectRef, val: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        if let Some(gen) = coro.payload::<PyGenerator>() {
-            gen.send(val, vm)
-        } else if let Some(coro) = coro.payload::<PyCoroutine>() {
-            coro.send(val, vm)
-        } else if vm.is_none(&val) {
-            objiter::call_next(vm, &coro)
-        } else {
-            vm.call_method(&coro, "send", vec![val])
+        let builtin_coro = match_class!(match coro {
+            ref g @ PyGenerator => Some(g.as_coro()),
+            ref c @ PyCoroutine => Some(c.as_coro()),
+            _ => None,
+        });
+        match builtin_coro {
+            Some(coro) => coro.send(val, vm),
+            None if vm.is_none(&val) => objiter::call_next(vm, &coro),
+            None => vm.call_method(&coro, "send", vec![val]),
         }
     }
 

vm/src/obj/mod.rs

@@ -8,6 +8,7 @@ pub mod objbytes;
 pub mod objclassmethod;
 pub mod objcode;
 pub mod objcomplex;
+pub mod objcoroinner;
 pub mod objcoroutine;
 pub mod objdict;
 pub mod objellipsis;

vm/src/obj/objcoroinner.rs

@@ -0,0 +1,84 @@
+use super::{objiter, objtype};
+use crate::exceptions::{self, PyBaseExceptionRef};
+use crate::frame::{ExecutionResult, FrameRef};
+use crate::pyobject::{PyObjectRef, PyResult};
+use crate::vm::VirtualMachine;
+
+use std::cell::Cell;
+
+#[derive(Debug)]
+pub struct Coro {
+    frame: FrameRef,
+    closed: Cell<bool>,
+}
+
+impl Coro {
+    pub fn new(frame: FrameRef) -> Self {
+        Coro {
+            frame,
+            closed: Cell::new(false),
+        }
+    }
+
+    fn maybe_close(&self, res: &PyResult<ExecutionResult>) {
+        match res {
+            Ok(ExecutionResult::Return(_)) | Err(_) => self.closed.set(true),
+            Ok(ExecutionResult::Yield(_)) => {}
+        }
+    }
+
+    pub fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        if self.closed.get() {
+            return Err(objiter::new_stop_iteration(vm));
+        }
+
+        self.frame.push_value(value.clone());
+
+        let result = vm.run_frame(self.frame.clone());
+        self.maybe_close(&result);
+        result?.into_result(vm)
+    }
+
+    pub fn throw(
+        &self,
+        exc_type: PyObjectRef,
+        exc_val: PyObjectRef,
+        exc_tb: PyObjectRef,
+        vm: &VirtualMachine,
+    ) -> PyResult {
+        if self.closed.get() {
+            return Err(exceptions::normalize(exc_type, exc_val, exc_tb, vm)?);
+        }
+        vm.frames.borrow_mut().push(self.frame.clone());
+        let result = self.frame.gen_throw(vm, exc_type, exc_val, exc_tb);
+        self.maybe_close(&result);
+        vm.frames.borrow_mut().pop();
+        result?.into_result(vm)
+    }
+
+    pub fn close(&self, vm: &VirtualMachine) -> PyResult<()> {
+        if self.closed.get() {
+            return Ok(());
+        }
+        vm.frames.borrow_mut().push(self.frame.clone());
+        let result = self.frame.gen_throw(
+            vm,
+            vm.ctx.exceptions.generator_exit.clone().into_object(),
+            vm.get_none(),
+            vm.get_none(),
+        );
+        vm.frames.borrow_mut().pop();
+        self.closed.set(true);
+        match result {
+            Ok(ExecutionResult::Yield(_)) => {
+                Err(vm.new_runtime_error("generator ignored GeneratorExit".to_owned()))
+            }
+            Err(e) if !is_gen_exit(&e, vm) => Err(e),
+            _ => Ok(()),
+        }
+    }
+}
+
+pub fn is_gen_exit(exc: &PyBaseExceptionRef, vm: &VirtualMachine) -> bool {
+    objtype::isinstance(exc, &vm.ctx.exceptions.generator_exit)
+}

vm/src/obj/objcoroutine.rs

@@ -1,20 +1,16 @@
-use super::objiter::new_stop_iteration;
-use super::objtype::{isinstance, PyClassRef};
-use crate::exceptions;
-use crate::frame::{ExecutionResult, FrameRef};
+use super::objcoroinner::Coro;
+use super::objtype::PyClassRef;
+use crate::frame::FrameRef;
 use crate::function::OptionalArg;
 use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
 use crate::vm::VirtualMachine;
 
-use std::cell::Cell;
-
 pub type PyCoroutineRef = PyRef<PyCoroutine>;
 
 #[pyclass(name = "coroutine")]
 #[derive(Debug)]
 pub struct PyCoroutine {
-    frame: FrameRef,
-    closed: Cell<bool>,
+    inner: Coro,
 }
 
 impl PyValue for PyCoroutine {
@@ -25,33 +21,20 @@ impl PyValue for PyCoroutine {
 
 #[pyimpl]
 impl PyCoroutine {
+    pub fn as_coro(&self) -> &Coro {
+        &self.inner
+    }
+
     pub fn new(frame: FrameRef, vm: &VirtualMachine) -> PyCoroutineRef {
         PyCoroutine {
-            frame,
-            closed: Cell::new(false),
+            inner: Coro::new(frame),
         }
         .into_ref(vm)
     }
 
-    // TODO: deduplicate this code with objgenerator
-    fn maybe_close(&self, res: &PyResult<ExecutionResult>) {
-        match res {
-            Ok(ExecutionResult::Return(_)) | Err(_) => self.closed.set(true),
-            Ok(ExecutionResult::Yield(_)) => {}
-        }
-    }
-
     #[pymethod]
-    pub(crate) fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        if self.closed.get() {
-            return Err(new_stop_iteration(vm));
-        }
-
-        self.frame.push_value(value.clone());
-
-        let result = vm.run_frame(self.frame.clone());
-        self.maybe_close(&result);
-        result?.into_result(vm)
+    fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        self.inner.send(value, vm)
     }
 
     #[pymethod]
@@ -62,45 +45,17 @@ impl PyCoroutine {
         exc_tb: OptionalArg,
         vm: &VirtualMachine,
     ) -> PyResult {
-        let exc_val = exc_val.unwrap_or_else(|| vm.get_none());
-        let exc_tb = exc_tb.unwrap_or_else(|| vm.get_none());
-        if self.closed.get() {
-            return Err(exceptions::normalize(exc_type, exc_val, exc_tb, vm)?);
-        }
-        vm.frames.borrow_mut().push(self.frame.clone());
-        let result = self.frame.gen_throw(vm, exc_type, exc_val, exc_tb);
-        self.maybe_close(&result);
-        vm.frames.borrow_mut().pop();
-        result?.into_result(vm)
+        self.inner.throw(
+            exc_type,
+            exc_val.unwrap_or_else(|| vm.get_none()),
+            exc_tb.unwrap_or_else(|| vm.get_none()),
+            vm,
+        )
     }
 
     #[pymethod]
     fn close(&self, vm: &VirtualMachine) -> PyResult<()> {
-        if self.closed.get() {
-            return Ok(());
-        }
-        vm.frames.borrow_mut().push(self.frame.clone());
-        let result = self.frame.gen_throw(
-            vm,
-            vm.ctx.exceptions.generator_exit.clone().into_object(),
-            vm.get_none(),
-            vm.get_none(),
-        );
-        vm.frames.borrow_mut().pop();
-        self.closed.set(true);
-        match result {
-            Ok(ExecutionResult::Yield(_)) => {
-                Err(vm.new_runtime_error("generator ignored GeneratorExit".to_owned()))
-            }
-            Err(e) => {
-                if isinstance(&e, &vm.ctx.exceptions.generator_exit) {
-                    Ok(())
-                } else {
-                    Err(e)
-                }
-            }
-            _ => Ok(()),
-        }
+        self.inner.close(vm)
     }
 
     #[pymethod(name = "__await__")]

vm/src/obj/objgenerator.rs

@@ -2,23 +2,19 @@
  * The mythical generator.
  */
 
-use super::objiter::new_stop_iteration;
-use super::objtype::{isinstance, PyClassRef};
-use crate::exceptions;
-use crate::frame::{ExecutionResult, FrameRef};
+use super::objcoroinner::Coro;
+use super::objtype::PyClassRef;
+use crate::frame::FrameRef;
 use crate::function::OptionalArg;
 use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
 use crate::vm::VirtualMachine;
 
-use std::cell::Cell;
-
 pub type PyGeneratorRef = PyRef<PyGenerator>;
 
 #[pyclass(name = "generator")]
 #[derive(Debug)]
 pub struct PyGenerator {
-    frame: FrameRef,
-    closed: Cell<bool>,
+    inner: Coro,
 }
 
 impl PyValue for PyGenerator {
@@ -29,19 +25,15 @@ impl PyValue for PyGenerator {
 
 #[pyimpl]
 impl PyGenerator {
-    pub fn new(frame: FrameRef, vm: &VirtualMachine) -> PyGeneratorRef {
-        PyGenerator {
-            frame,
-            closed: Cell::new(false),
-        }
-        .into_ref(vm)
+    pub fn as_coro(&self) -> &Coro {
+        &self.inner
     }
 
-    fn maybe_close(&self, res: &PyResult<ExecutionResult>) {
-        match res {
-            Ok(ExecutionResult::Return(_)) | Err(_) => self.closed.set(true),
-            Ok(ExecutionResult::Yield(_)) => {}
+    pub fn new(frame: FrameRef, vm: &VirtualMachine) -> PyGeneratorRef {
+        PyGenerator {
+            inner: Coro::new(frame),
         }
+        .into_ref(vm)
     }
 
     #[pymethod(name = "__iter__")]
@@ -55,16 +47,8 @@ impl PyGenerator {
     }
 
     #[pymethod]
-    pub(crate) fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        if self.closed.get() {
-            return Err(new_stop_iteration(vm));
-        }
-
-        self.frame.push_value(value.clone());
-
-        let result = vm.run_frame(self.frame.clone());
-        self.maybe_close(&result);
-        result?.into_result(vm)
+    fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        self.inner.send(value, vm)
     }
 
     #[pymethod]
@@ -75,45 +59,17 @@ impl PyGenerator {
         exc_tb: OptionalArg,
         vm: &VirtualMachine,
     ) -> PyResult {
-        let exc_val = exc_val.unwrap_or_else(|| vm.get_none());
-        let exc_tb = exc_tb.unwrap_or_else(|| vm.get_none());
-        if self.closed.get() {
-            return Err(exceptions::normalize(exc_type, exc_val, exc_tb, vm)?);
-        }
-        vm.frames.borrow_mut().push(self.frame.clone());
-        let result = self.frame.gen_throw(vm, exc_type, exc_val, exc_tb);
-        self.maybe_close(&result);
-        vm.frames.borrow_mut().pop();
-        result?.into_result(vm)
+        self.inner.throw(
+            exc_type,
+            exc_val.unwrap_or_else(|| vm.get_none()),
+            exc_tb.unwrap_or_else(|| vm.get_none()),
+            vm,
+        )
     }
 
     #[pymethod]
     fn close(&self, vm: &VirtualMachine) -> PyResult<()> {
-        if self.closed.get() {
-            return Ok(());
-        }
-        vm.frames.borrow_mut().push(self.frame.clone());
-        let result = self.frame.gen_throw(
-            vm,
-            vm.ctx.exceptions.generator_exit.clone().into_object(),
-            vm.get_none(),
-            vm.get_none(),
-        );
-        vm.frames.borrow_mut().pop();
-        self.closed.set(true);
-        match result {
-            Ok(ExecutionResult::Yield(_)) => {
-                Err(vm.new_runtime_error("generator ignored GeneratorExit".to_owned()))
-            }
-            Err(e) => {
-                if isinstance(&e, &vm.ctx.exceptions.generator_exit) {
-                    Ok(())
-                } else {
-                    Err(e)
-                }
-            }
-            _ => Ok(()),
-        }
+        self.inner.close(vm)
     }
 }
 

vm/src/pyobject.rs

@@ -640,11 +640,9 @@ impl PyObject<dyn PyObjectPayload> {
     /// another downcast can be attempted without unnecessary cloning.
     pub fn downcast<T: PyObjectPayload>(self: Rc<Self>) -> Result<PyRef<T>, PyObjectRef> {
         if self.payload_is::<T>() {
-            Ok({
-                PyRef {
-                    obj: self,
-                    _payload: PhantomData,
-                }
+            Ok(PyRef {
+                obj: self,
+                _payload: PhantomData,
             })
         } else {
             Err(self)