objgenerator: Keep exception stack within generator object, like value stack.

pfalcon · pfalcon · commit c0abc28aa121 · 2014-03-22T17:55:42.000+02:00
This is required to properly handle exceptions across yields.
diff --git a/py/bc.h b/py/bc.h
@@ -4,6 +4,16 @@ typedef enum {
     MP_VM_RETURN_EXCEPTION,
 } mp_vm_return_kind_t;
 
+// Exception stack entry
+typedef struct _mp_exc_stack {
+    const byte *handler;
+    // bit 0 is saved currently_in_except_block value
+    machine_uint_t val_sp;
+    // We might only have 2 interesting cases here: SETUP_EXCEPT & SETUP_FINALLY,
+    // consider storing it in bit 1 of val_sp. TODO: SETUP_WITH?
+    byte opcode;
+} mp_exc_stack;
+
 mp_vm_return_kind_t mp_execute_byte_code(const byte *code, const mp_obj_t *args, uint n_args, const mp_obj_t *args2, uint n_args2, uint n_state, mp_obj_t *ret);
-mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **ip_in_out, mp_obj_t *fastn, mp_obj_t **sp_in_out);
+mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **ip_in_out, mp_obj_t *fastn, mp_obj_t **sp_in_out, mp_exc_stack *exc_stack, mp_exc_stack **exc_sp_in_out);
 void mp_byte_code_print(const byte *code, int len);
diff --git 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) {
@@ -80,7 +82,8 @@ 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);
     switch (vm_return_kind) {
         case MP_VM_RETURN_NORMAL:
             // Explicitly mark generator as completed. If we don't do this,
@@ -137,11 +140,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)
diff --git a/py/vm.c b/py/vm.c
@@ -17,16 +17,6 @@
 // top element.
 // Exception stack also grows up, top element is also pointed at.
 
-// Exception stack entry
-typedef struct _mp_exc_stack {
-    const byte *handler;
-    // bit 0 is saved currently_in_except_block value
-    machine_uint_t val_sp;
-    // We might only have 2 interesting cases here: SETUP_EXCEPT & SETUP_FINALLY,
-    // consider storing it in bit 1 of val_sp. TODO: SETUP_WITH?
-    byte opcode;
-} mp_exc_stack;
-
 // Exception stack unwind reasons (WHY_* in CPython-speak)
 // TODO perhaps compress this to RETURN=0, JUMP>0, with number of unwinds
 // left to do encoded in the JUMP number
@@ -89,8 +79,10 @@ mp_vm_return_kind_t mp_execute_byte_code(const byte *code, const mp_obj_t *args,
         }
     }
 
+    mp_exc_stack exc_stack[4];
+    mp_exc_stack *exc_sp = &exc_stack[0] - 1;
     // execute the byte code
-    mp_vm_return_kind_t vm_return_kind = mp_execute_byte_code_2(code, &ip, &state[n_state - 1], &sp);
+    mp_vm_return_kind_t vm_return_kind = mp_execute_byte_code_2(code, &ip, &state[n_state - 1], &sp, exc_stack, &exc_sp);
 
     switch (vm_return_kind) {
         case MP_VM_RETURN_NORMAL:
@@ -113,7 +105,10 @@ mp_vm_return_kind_t mp_execute_byte_code(const byte *code, const mp_obj_t *args,
 //  MP_VM_RETURN_NORMAL, sp valid, return value in *sp
 //  MP_VM_RETURN_YIELD, ip, sp valid, yielded value in *sp
 //  MP_VM_RETURN_EXCEPTION, exception in fastn[0]
-mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **ip_in_out, mp_obj_t *fastn, mp_obj_t **sp_in_out) {
+mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **ip_in_out,
+                                           mp_obj_t *fastn, mp_obj_t **sp_in_out,
+                                           mp_exc_stack *exc_stack,
+                                           mp_exc_stack **exc_sp_in_out) {
     // careful: be sure to declare volatile any variables read in the exception handler (written is ok, I think)
 
     const byte *ip = *ip_in_out;
@@ -124,8 +119,7 @@ mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **i
     nlr_buf_t nlr;
 
     volatile machine_uint_t currently_in_except_block = 0; // 0 or 1, to detect nested exceptions
-    mp_exc_stack exc_stack[4];
-    mp_exc_stack *volatile exc_sp = &exc_stack[0] - 1; // stack grows up, exc_sp points to top of stack
+    mp_exc_stack *volatile exc_sp = *exc_sp_in_out; // stack grows up, exc_sp points to top of stack
     const byte *volatile save_ip = ip; // this is so we can access ip in the exception handler without making ip volatile (which means the compiler can't keep it in a register in the main loop)
 
     // outer exception handling loop
@@ -434,7 +428,7 @@ mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **i
                     // matched against: SETUP_EXCEPT, SETUP_FINALLY, SETUP_WITH
                     case MP_BC_POP_BLOCK:
                         // we are exiting an exception handler, so pop the last one of the exception-stack
-                        assert(exc_sp >= &exc_stack[0]);
+                        assert(exc_sp >= exc_stack);
                         currently_in_except_block = (exc_sp->val_sp & 1); // restore previous state
                         exc_sp--; // pop back to previous exception handler
                         break;
@@ -443,7 +437,7 @@ mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **i
                     case MP_BC_POP_EXCEPT:
                         // TODO need to work out how blocks work etc
                         // pops block, checks it's an exception block, and restores the stack, saving the 3 exception values to local threadstate
-                        assert(exc_sp >= &exc_stack[0]);
+                        assert(exc_sp >= exc_stack);
                         assert(currently_in_except_block);
                         //sp = (mp_obj_t*)(*exc_sp--);
                         //exc_sp--; // discard ip
@@ -592,7 +586,7 @@ mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **i
                         }
                         nlr_pop();
                         *sp_in_out = sp;
-                        assert(exc_sp == &exc_stack[0] - 1);
+                        assert(exc_sp == exc_stack - 1);
                         return MP_VM_RETURN_NORMAL;
 
                     case MP_BC_RAISE_VARARGS:
@@ -605,6 +599,7 @@ mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **i
                         nlr_pop();
                         *ip_in_out = ip;
                         *sp_in_out = sp;
+                        *exc_sp_in_out = exc_sp;
                         return MP_VM_RETURN_YIELD;
 
                     case MP_BC_IMPORT_NAME:
@@ -654,7 +649,7 @@ mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **i
             while (currently_in_except_block) {
                 // nested exception
 
-                assert(exc_sp >= &exc_stack[0]);
+                assert(exc_sp >= exc_stack);
 
                 // TODO make a proper message for nested exception
                 // at the moment we are just raising the very last exception (the one that caused the nested exception)
@@ -664,7 +659,7 @@ mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **i
                 exc_sp--; // pop back to previous exception handler
             }
 
-            if (exc_sp >= &exc_stack[0]) {
+            if (exc_sp >= exc_stack) {
                 // set flag to indicate that we are now handling an exception
                 currently_in_except_block = 1;
 
diff --git a/tests/basics/generator-exc.py b/tests/basics/generator-exc.py
@@ -0,0 +1,11 @@
+# Test proper handling of exceptions within generator across yield
+def gen():
+    try:
+        yield 1
+        raise ValueError
+    except ValueError:
+        print("Caught")
+    yield 2
+
+for i in gen():
+    print(i)
