1 files changed, 50 insertions, 10 deletions

diff --git a/src/vm.c b/src/vm.c
index 4b121cecb..9b9cb3869 100644
--- a/src/vm.c
+++ b/src/vm.c
@@ -302,6 +302,7 @@ ecall(mrb_state *mrb, int i)
   mrb_value *self = mrb->c->stack;
   struct RObject *exc;
   ptrdiff_t cioff;
+  int ai = mrb_gc_arena_save(mrb);
 
   if (i<0) return;
   if (ci - mrb->c->cibase > MRB_FUNCALL_DEPTH_MAX) {
@@ -321,9 +322,13 @@ ecall(mrb_state *mrb, int i)
   ci->target_class = p->target_class;
   mrb->c->stack = mrb->c->stack + ci[-1].nregs;
   exc = mrb->exc; mrb->exc = 0;
+  if (exc) {
+    mrb_gc_protect(mrb, mrb_obj_value(exc));
+  }
   mrb_run(mrb, p, *self);
   mrb->c->ci = mrb->c->cibase + cioff;
   if (!mrb->exc) mrb->exc = exc;
+  mrb_gc_arena_restore(mrb, ai);
 }
 
 #ifndef MRB_FUNCALL_ARGC_MAX
@@ -755,6 +760,19 @@ mrb_yield_cont(mrb_state *mrb, mrb_value b, mrb_value self, mrb_int argc, const
   return mrb_exec_irep(mrb, self, p);
 }
 
+static struct RBreak*
+break_new(mrb_state *mrb, struct RProc *p, mrb_value val)
+{
+  struct RBreak *brk;
+
+  brk = (struct RBreak*)mrb_obj_alloc(mrb, MRB_TT_BREAK, NULL);
+  brk->iv = NULL;
+  brk->proc = p;
+  brk->val = val;
+
+  return brk;
+}
+
 typedef enum {
   LOCALJUMP_ERROR_RETURN = 0,
   LOCALJUMP_ERROR_BREAK = 1,
@@ -862,7 +880,12 @@ mrb_vm_run(mrb_state *mrb, struct RProc *proc, mrb_value self, unsigned int stac
   if (c->ci - c->cibase > cioff) {
     c->ci = c->cibase + cioff;
   }
-  mrb->c = c;
+  if (mrb->c != c) {
+    if (mrb->c->fib) {
+      mrb_write_barrier(mrb, (struct RBasic*)mrb->c->fib);
+    }
+    mrb->c = c;
+  }
   return result;
 }
 
@@ -910,6 +933,8 @@ RETRY_TRY_BLOCK:
 
   if (exc_catched) {
     exc_catched = FALSE;
+    if (mrb->exc && mrb->exc->tt == MRB_TT_BREAK)
+      goto L_BREAK;
     goto L_RAISE;
   }
   mrb->jmp = &c_jmp;
@@ -1185,7 +1210,9 @@ RETRY_TRY_BLOCK:
 
     CASE(OP_RAISE) {
       /* A      raise(R(A)) */
-      mrb_exc_set(mrb, regs[GETARG_A(i)]);
+      int a = GETARG_A(i);
+
+      mrb_exc_set(mrb, regs[a]);
       goto L_RAISE;
     }
 
@@ -1434,14 +1461,13 @@ RETRY_TRY_BLOCK:
         pool = irep->pool;
         syms = irep->syms;
         ci->nregs = irep->nregs;
+        stack_extend(mrb, irep->nregs);
         if (ci->argc < 0) {
           if (irep->nregs > 3) {
-            stack_extend(mrb, irep->nregs);
             stack_clear(regs+3, irep->nregs-3);
           }
         }
         else if (ci->argc+2 < irep->nregs) {
-          stack_extend(mrb, irep->nregs);
           stack_clear(regs+ci->argc+2, irep->nregs-ci->argc-2);
         }
         if (m->env) {
@@ -1813,6 +1839,9 @@ RETRY_TRY_BLOCK:
               else {
                 struct mrb_context *c = mrb->c;
 
+                if (c->fib) {
+                  mrb_write_barrier(mrb, (struct RBasic*)c->fib);
+                }
                 mrb->c = c->prev;
                 c->prev = NULL;
                 goto L_RAISE;
@@ -1837,12 +1866,14 @@ RETRY_TRY_BLOCK:
         if (ci != ci0) {
           mrb->c->stack = ci[1].stackent;
         }
+        stack_extend(mrb, irep->nregs);
         pc = mrb->c->rescue[--ci->ridx];
       }
       else {
         int acc;
-        mrb_value v = regs[GETARG_A(i)];
+        mrb_value v;
 
+        v = regs[GETARG_A(i)];
         mrb_gc_protect(mrb, v);
         switch (GETARG_B(i)) {
         case OP_R_RETURN:
@@ -1857,7 +1888,7 @@ RETRY_TRY_BLOCK:
             }
             
             ce = mrb->c->cibase + e->cioff;
-            while (--ci > ce) {
+            while (ci >= ce) {
               if (ci->env) {
                 mrb_env_unshare(mrb, ci->env);
               }
@@ -1865,6 +1896,7 @@ RETRY_TRY_BLOCK:
                 localjump_error(mrb, LOCALJUMP_ERROR_RETURN);
                 goto L_RAISE;
               }
+              ci--;
             }
             if (ce == mrb->c->cibase) {
               localjump_error(mrb, LOCALJUMP_ERROR_RETURN);
@@ -1927,19 +1959,27 @@ RETRY_TRY_BLOCK:
             }
             ARENA_RESTORE(mrb, ai);
             mrb->c->vmexec = FALSE;
+            mrb->exc = (struct RObject*)break_new(mrb, proc, v);
             mrb->jmp = prev_jmp;
-            return v;
+            MRB_THROW(prev_jmp);
+          }
+          if (FALSE) {
+          L_BREAK:
+            v = ((struct RBreak*)mrb->exc)->val;
+            proc = ((struct RBreak*)mrb->exc)->proc;
+            mrb->exc = NULL;
+            ci = mrb->c->ci;
           }
           mrb->c->stack = ci->stackent;
           mrb->c->ci = mrb->c->cibase + proc->env->cioff + 1;
           while (ci > mrb->c->ci) {
+            if (ci->env) {
+              mrb_env_unshare(mrb, ci->env);
+            }
             if (ci[-1].acc == CI_ACC_SKIP) {
               mrb->c->ci = ci;
               goto L_BREAK_ERROR;
             }
-            if (ci->env) {
-              mrb_env_unshare(mrb, ci->env);
-            }
             ci--;
           }
           break;