From ad2e626e7a937fb4de9a1647aa58ea45aa3936ec Mon Sep 17 00:00:00 2001
From: dearblue <dearblue@users.osdn.me>
Date: Fri, 31 May 2024 22:03:15 +0900
Subject: [PATCH] Delegate the care of a directly given block from `cipop()` to
 `cipush()`

Outlines:
  - Removed `mrb_callinfo::blk`
  - Added `mrb_callinfo::flags`
  - Added `MRB_CI_COMPANION_BLOCK` flag
---
 include/mruby.h          |  2 +-
 include/mruby/internal.h |  3 +++
 src/vm.c                 | 30 ++++++++++++++++++++++--------
 3 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/include/mruby.h b/include/mruby.h
index 3b4a1fc72..490f344ac 100644
--- a/include/mruby.h
+++ b/include/mruby.h
@@ -176,9 +176,9 @@ typedef struct {
   uint8_t n:4;                  /* (15=*) c=n|nk<<4 */
   uint8_t nk:4;                 /* (15=*) */
   uint8_t cci;                  /* called from C function */
+  uint8_t flags;                /* MRB_CI_COMPANION_BLOCK or zero */
   mrb_sym mid;
   const struct RProc *proc;
-  struct RProc *blk;
   mrb_value *stack;
   const mrb_code *pc;           /* current address on iseq of this proc */
   union {
diff --git a/include/mruby/internal.h b/include/mruby/internal.h
index 9f8453eaa..d8e81eafe 100644
--- a/include/mruby/internal.h
+++ b/include/mruby/internal.h
@@ -195,6 +195,9 @@ mrb_value mrb_obj_instance_eval(mrb_state*, mrb_value);
 mrb_value mrb_mod_module_eval(mrb_state*, mrb_value);
 mrb_value mrb_f_send(mrb_state *mrb, mrb_value self);
 
+/* mrb_callinfo::flags */
+#define MRB_CI_COMPANION_BLOCK  0x01    /* it means `method { ... }`, not `method(&blk)` */
+
 #ifdef MRB_USE_BIGINT
 mrb_value mrb_bint_new_int(mrb_state *mrb, mrb_int x);
 #ifdef MRB_INT64
diff --git a/src/vm.c b/src/vm.c
index 1b94ee5e1..0cc858a71 100644
--- a/src/vm.c
+++ b/src/vm.c
@@ -360,9 +360,15 @@ cipush(mrb_state *mrb, mrb_int push_stacks, uint8_t cci, struct RClass *target_c
     c->ciend = c->cibase + size * 2;
   }
   ci = ++c->ci;
+  ci->flags = 0;
+  if (blk && (blk->flags & (MRB_PROC_CFUNC_FL | MRB_PROC_ENVSET | MRB_PROC_ORPHAN)) == MRB_PROC_ENVSET &&
+      blk->e.env == ci[-1].u.env) {
+    mrb_assert(blk->color != MRB_GC_RED); // no exist red object with env set
+    ci->flags = MRB_CI_COMPANION_BLOCK;
+    ((struct RProc*)blk)->flags |= MRB_PROC_ORPHAN;
+  }
   ci->mid = mid;
   CI_PROC_SET(ci, proc);
-  ci->blk = blk;
   ci->stack = ci[-1].stack + push_stacks;
   ci->n = argc & 0xf;
   ci->nk = (argc>>4) & 0xf;
@@ -474,11 +480,6 @@ cipop(mrb_state *mrb)
   struct REnv *env = CI_ENV(ci);
 
   ci_env_set(ci, NULL); // make possible to free env by GC if not needed
-  struct RProc *b = ci->blk;
-  if (b && !mrb_object_dead_p(mrb, (struct RBasic*)b) && b->tt == MRB_TT_PROC &&
-      !MRB_PROC_STRICT_P(b) && MRB_PROC_ENV(b) == CI_ENV(&ci[-1])) {
-    b->flags |= MRB_PROC_ORPHAN;
-  }
   if (env && !mrb_env_unshare(mrb, env, TRUE)) {
     c->ci--; // exceptions are handled at the method caller; see #3087
     mrb_exc_raise(mrb, mrb_obj_value(mrb->nomem_err));
@@ -2305,14 +2306,27 @@ RETRY_TRY_BLOCK:
       }
 
       if (MRB_PROC_STRICT_P(proc)) goto NORMAL_RETURN;
-      if (MRB_PROC_ORPHAN_P(proc) || !MRB_PROC_ENV_P(proc) || !MRB_ENV_ONSTACK_P(MRB_PROC_ENV(proc))) {
+      if (!MRB_PROC_ENV_P(proc)) {
       L_BREAK_ERROR:
         RAISE_LIT(mrb, E_LOCALJUMP_ERROR, "break from proc-closure");
       }
       else {
         struct REnv *e = MRB_PROC_ENV(proc);
 
-        if (e->cxt != mrb->c) {
+        if (!MRB_ENV_ONSTACK_P(e) || e->cxt != mrb->c) {
+          goto L_BREAK_ERROR;
+        }
+
+        mrb_callinfo *birth_ci = mrb->c->ci - 1;
+        for (; birth_ci >= mrb->c->cibase; birth_ci--) {
+          if (e == birth_ci->u.env) {
+            if (!(birth_ci[1].flags & MRB_CI_COMPANION_BLOCK)) {
+              goto L_BREAK_ERROR;
+            }
+            break;
+          }
+        }
+        if (birth_ci < mrb->c->cibase) {
           goto L_BREAK_ERROR;
         }
       }