Allow `break` from a block called by `mrb_yield`; close #3359

This means #3701 is now OK to merge.

4 files changed, 39 insertions, 6 deletions

diff --git a/include/mruby/error.h b/include/mruby/error.h
index bb67e7bd8..0a262550e 100644
--- a/include/mruby/error.h
+++ b/include/mruby/error.h
@@ -32,6 +32,13 @@ MRB_API mrb_noreturn void mrb_no_method_error(mrb_state *mrb, mrb_sym id, mrb_va
 /* declaration for fail method */
 MRB_API mrb_value mrb_f_raise(mrb_state*, mrb_value);
 
+struct RBreak {
+  MRB_OBJECT_HEADER;
+  struct iv_tbl *iv;
+  struct RProc *proc;
+  mrb_value val;
+};
+
 /**
  * Protect
  *
diff --git a/include/mruby/value.h b/include/mruby/value.h
index a206be549..98c68d657 100644
--- a/include/mruby/value.h
+++ b/include/mruby/value.h
@@ -116,7 +116,8 @@ enum mrb_vtype {
   MRB_TT_DATA,        /*  21 */
   MRB_TT_FIBER,       /*  22 */
   MRB_TT_ISTRUCT,     /*  23 */
-  MRB_TT_MAXDEFINE    /*  24 */
+  MRB_TT_BREAK,       /*  24 */
+  MRB_TT_MAXDEFINE    /*  25 */
 };
 
 #include <mruby/object.h>
diff --git a/src/gc.c b/src/gc.c
index 781059d13..f74a7f31d 100644
--- a/src/gc.c
+++ b/src/gc.c
@@ -112,6 +112,7 @@ typedef struct {
     struct RProc proc;
     struct REnv env;
     struct RException exc;
+    struct RBreak brk;
 #ifdef MRB_WORD_BOXING
     struct RFloat floatv;
     struct RCptr cptr;
diff --git a/src/vm.c b/src/vm.c
index 9607d3c32..9b9cb3869 100644
--- a/src/vm.c
+++ b/src/vm.c
@@ -760,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,
@@ -920,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;
@@ -1856,8 +1871,9 @@ RETRY_TRY_BLOCK:
       }
       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:
@@ -1943,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;