1 files changed, 16 insertions, 5 deletions
diff --git a/py/vm.c b/py/vm.c
index f6aa74348d..75093d2401 100644
--- a/py/vm.c
+++ b/py/vm.c
@@ -103,7 +103,8 @@ typedef enum {
     currently_in_except_block = MP_TAGPTR_TAG(exc_sp->val_sp); /* restore previous state */ \
     exc_sp--; /* pop back to previous exception handler */
 
-mp_vm_return_kind_t mp_execute_bytecode(const byte *code, const mp_obj_t *args, uint n_args, const mp_obj_t *args2, uint n_args2, mp_obj_t *ret) {
+mp_vm_return_kind_t mp_execute_bytecode(const byte *code, const mp_obj_t *args, uint n_args,
+                                                          const mp_obj_t *args2, uint n_args2, mp_obj_t *ret) {
     const byte *ip = code;
 
     // get code info size, and skip line number table
@@ -157,6 +158,13 @@ mp_vm_return_kind_t mp_execute_bytecode(const byte *code, const mp_obj_t *args,
     mp_vm_return_kind_t vm_return_kind = mp_execute_bytecode2(code, &ip, &state[n_state - 1], &sp, exc_stack, &exc_sp, MP_OBJ_NULL);
 
 #if DETECT_VM_STACK_OVERFLOW
+    if (vm_return_kind == MP_VM_RETURN_NORMAL) {
+        if (sp < state) {
+            printf("VM stack underflow: " INT_FMT "\n", sp - state);
+            assert(0);
+        }
+    }
+
     // We can't check the case when an exception is returned in state[n_state - 1]
     // and there are no arguments, because in this case our detection slot may have
     // been overwritten by the returned exception (which is allowed).
@@ -170,7 +178,7 @@ mp_vm_return_kind_t mp_execute_bytecode(const byte *code, const mp_obj_t *args,
             }
         }
         if (overflow) {
-            printf("VM stack overflow state=%p n_state+1=%u\n", state, n_state);
+            printf("VM stack overflow state=%p n_state+1=" UINT_FMT "\n", state, n_state);
             assert(0);
         }
     }
@@ -312,7 +320,7 @@ dispatch_loop:
 
                 ENTRY(MP_BC_LOAD_CONST_INT): {
                     DECODE_QSTR;
-                    PUSH(mp_obj_new_int_from_qstr(qst));
+                    PUSH(mp_load_const_int(qst));
                     DISPATCH();
                 }
 
@@ -618,10 +626,10 @@ dispatch_loop:
                 ENTRY(MP_BC_UNWIND_JUMP):
                     DECODE_SLABEL;
                     PUSH((void*)(ip + unum)); // push destination ip for jump
-                    PUSH((void*)(machine_uint_t)(*ip)); // push number of exception handlers to unwind
+                    PUSH((void*)(machine_uint_t)(*ip)); // push number of exception handlers to unwind (0x80 bit set if we also need to pop stack)
 unwind_jump:
                     unum = (machine_uint_t)POP(); // get number of exception handlers to unwind
-                    while (unum > 0) {
+                    while ((unum & 0x7f) > 0) {
                         unum -= 1;
                         assert(exc_sp >= exc_stack);
                         if (exc_sp->opcode == MP_BC_SETUP_FINALLY || exc_sp->opcode == MP_BC_SETUP_WITH) {
@@ -638,6 +646,9 @@ unwind_jump:
                         exc_sp--;
                     }
                     ip = (const byte*)POP(); // pop destination ip for jump
+                    if (unum != 0) {
+                        sp--;
+                    }
                     DISPATCH();
 
                 // matched against: POP_BLOCK or POP_EXCEPT (anything else?)