1 files changed, 28 insertions, 7 deletions

diff --git a/py/objgenerator.c b/py/objgenerator.c
index 8c4bb595f1..2db04ad742 100644
--- a/py/objgenerator.c
+++ b/py/objgenerator.c
@@ -56,8 +56,10 @@ typedef struct _mp_obj_gen_instance_t {
     const byte *code_info;
     const byte *ip;
     mp_obj_t *sp;
+    mp_exc_stack *exc_sp;
     uint n_state;
-    mp_obj_t state[];
+    mp_obj_t state[0];          // Variable-length
+    mp_exc_stack exc_state[0];  // Variable-length
 } mp_obj_gen_instance_t;
 
 void gen_instance_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in, mp_print_kind_t kind) {
@@ -68,7 +70,7 @@ mp_obj_t gen_instance_getiter(mp_obj_t self_in) {
     return self_in;
 }
 
-STATIC mp_obj_t gen_next_send(mp_obj_t self_in, mp_obj_t send_value) {
+STATIC mp_obj_t gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw_value) {
     mp_obj_gen_instance_t *self = self_in;
     if (self->ip == 0) {
         return mp_const_stop_iteration;
@@ -80,7 +82,9 @@ STATIC mp_obj_t gen_next_send(mp_obj_t self_in, mp_obj_t send_value) {
     } else {
         *self->sp = send_value;
     }
-    mp_vm_return_kind_t vm_return_kind = mp_execute_byte_code_2(self->code_info, &self->ip, &self->state[self->n_state - 1], &self->sp);
+    mp_vm_return_kind_t vm_return_kind = mp_execute_byte_code_2(self->code_info, &self->ip,
+        &self->state[self->n_state - 1], &self->sp, (mp_exc_stack*)(self->state + self->n_state),
+        &self->exc_sp, throw_value);
     switch (vm_return_kind) {
         case MP_VM_RETURN_NORMAL:
             // Explicitly mark generator as completed. If we don't do this,
@@ -100,19 +104,21 @@ STATIC mp_obj_t gen_next_send(mp_obj_t self_in, mp_obj_t send_value) {
             return *self->sp;
 
         case MP_VM_RETURN_EXCEPTION:
+            self->ip = 0;
+            nlr_jump(self->state[self->n_state - 1]);
+
         default:
-            // TODO
             assert(0);
             return mp_const_none;
     }
 }
 
 mp_obj_t gen_instance_iternext(mp_obj_t self_in) {
-    return gen_next_send(self_in, mp_const_none);
+    return gen_resume(self_in, mp_const_none, MP_OBJ_NULL);
 }
 
 STATIC mp_obj_t gen_instance_send(mp_obj_t self_in, mp_obj_t send_value) {
-    mp_obj_t ret = gen_next_send(self_in, send_value);
+    mp_obj_t ret = gen_resume(self_in, send_value, MP_OBJ_NULL);
     if (ret == mp_const_stop_iteration) {
         nlr_jump(mp_obj_new_exception(&mp_type_StopIteration));
     } else {
@@ -122,8 +128,21 @@ STATIC mp_obj_t gen_instance_send(mp_obj_t self_in, mp_obj_t send_value) {
 
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(gen_instance_send_obj, gen_instance_send);
 
+STATIC mp_obj_t gen_instance_throw(uint n_args, const mp_obj_t *args) {
+    mp_obj_t ret = gen_resume(args[0], mp_const_none, n_args == 2 ? args[1] : args[2]);
+    if (ret == mp_const_stop_iteration) {
+        nlr_jump(mp_obj_new_exception(&mp_type_StopIteration));
+    } else {
+        return ret;
+    }
+}
+
+STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(gen_instance_throw_obj, 2, 4, gen_instance_throw);
+
+
 STATIC const mp_method_t gen_type_methods[] = {
     { "send", &gen_instance_send_obj },
+    { "throw", &gen_instance_throw_obj },
     { NULL, NULL }, // end-of-list sentinel
 };
 
@@ -137,11 +156,13 @@ const mp_obj_type_t gen_instance_type = {
 };
 
 mp_obj_t mp_obj_new_gen_instance(const byte *bytecode, uint n_state, int n_args, const mp_obj_t *args) {
-    mp_obj_gen_instance_t *o = m_new_obj_var(mp_obj_gen_instance_t, mp_obj_t, n_state);
+    // TODO: 4 is hardcoded number from vm.c, calc exc stack size instead.
+    mp_obj_gen_instance_t *o = m_new_obj_var(mp_obj_gen_instance_t, byte, n_state * sizeof(mp_obj_t) + 4 * sizeof(mp_exc_stack));
     o->base.type = &gen_instance_type;
     o->code_info = bytecode;
     o->ip = bytecode;
     o->sp = &o->state[0] - 1; // sp points to top of stack, which starts off 1 below the state
+    o->exc_sp = (mp_exc_stack*)(o->state + n_state) - 1;
     o->n_state = n_state;
 
     // copy args to end of state array, in reverse (that's how mp_execute_byte_code_2 needs it)