From c654123561d679c190b0fd05848576bd1bf41ec3 Mon Sep 17 00:00:00 2001
From: dearblue <dearblue@users.osdn.me>
Date: Sun, 10 Nov 2024 18:39:00 +0900
Subject: [PATCH] Distinguish the call frame of the generator with
 `OP_RETURN_BLK`

When multiple identical proc objects are placed on the call stack, it is not possible to distinguish where to `return`.
Therefore, use env object comparisons to do this.

fixed #6411
---
 src/vm.c           |  8 +++--
 test/t/bs_block.rb | 86 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+), 3 deletions(-)

diff --git a/src/vm.c b/src/vm.c
index dac658001..53d8f3329 100644
--- a/src/vm.c
+++ b/src/vm.c
@@ -233,11 +233,12 @@ uvenv(mrb_state *mrb, mrb_int up)
 }
 
 static inline const struct RProc*
-top_proc(mrb_state *mrb, const struct RProc *proc)
+top_proc(mrb_state *mrb, const struct RProc *proc, const struct REnv **envp)
 {
   while (proc->upper) {
     if (MRB_PROC_SCOPE_P(proc) || MRB_PROC_STRICT_P(proc))
       return proc;
+    *envp = proc->e.env;
     proc = proc->upper;
   }
   return proc;
@@ -2294,11 +2295,12 @@ RETRY_TRY_BLOCK:
         goto NORMAL_RETURN;
       }
 
-      const struct RProc *dst = top_proc(mrb, ci->proc);
+      const struct REnv *env = ci->u.env;
+      const struct RProc *dst = top_proc(mrb, ci->proc, &env);
       if (!MRB_PROC_ENV_P(dst) || dst->e.env->cxt == mrb->c) {
         /* check jump destination */
         for (ptrdiff_t i = ci - mrb->c->cibase; i >= 0; i--, ci--) {
-          if (ci->proc == dst) {
+          if (ci->u.env == env) {
             goto L_UNWINDING;
           }
         }
diff --git a/test/t/bs_block.rb b/test/t/bs_block.rb
index 08580d58a..f3ead4f2a 100644
--- a/test/t/bs_block.rb
+++ b/test/t/bs_block.rb
@@ -532,3 +532,89 @@ assert('BS Block 39') do
     }
   end
 end
+
+assert('BS Block 40 (https://github.com/mruby/mruby/issues/6411)') do
+  assert_equal "GOOD" do
+    Object.new.instance_eval do
+      def test(&b)
+        if b
+          b.call
+        else
+          test { return "GOOD" }
+        end
+        "BAD"
+      end
+
+      test
+    end
+  end
+
+  assert_equal "GOOD" do
+    Object.new.instance_eval do
+      # since Kernel#proc is defined in proc-ext
+      def make_proc(&b)
+        b
+      end
+
+      def chocolate(&b)
+        biscuit(&b)
+      end
+
+      def biscuit(&b)
+        if b
+          b.call
+        else
+          b = make_proc { return "GOOD" }
+          chocolate(&b)
+        end
+        "BAD"
+      end
+
+      biscuit
+    end
+  end
+
+  assert_equal [0, 1, 2, 3] do
+    Object.new.instance_eval do
+      def test(a = [], &b)
+        if b
+          b.call
+        else
+          if a.empty?
+            a << 0
+            test(a)
+          else
+            a << 1
+            test(a) { return 1 }
+          end
+          a << 2
+        end
+        a << 3
+      end
+
+      test
+    end
+  end
+
+  assert_equal [0, 1, 3, 2, 3, 2, 3] do
+    Object.new.instance_eval do
+      def test(a = [], &b)
+        if b
+          b.call
+        else
+          if a.empty?
+            a << 0
+            test(a)
+          else
+            a << 1
+            test(a, &-> { return 1 })
+          end
+          a << 2
+        end
+        a << 3
+      end
+
+      test
+    end
+  end
+end