Merge pull request #2716 from kou/fix-ensure-context-on-break-and-return

Fix ensure with yield context on break and return

src/vm.c

@@ -1514,6 +1514,7 @@ RETRY_TRY_BLOCK:
               localjump_error(mrb, LOCALJUMP_ERROR_RETURN);
               goto L_RAISE;
             }
+            mrb->c->stack = mrb->c->ci->stackent;
             mrb->c->ci = ci;
             break;
           }
@@ -1548,6 +1549,7 @@ RETRY_TRY_BLOCK:
             c->prev = 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[-1].acc == CI_ACC_SKIP) {

test/t/ensure.rb

+##
+# ensure Test
+
+assert('ensure - context - yield') do
+  class EnsureYieldBreak
+    attr_reader :ensure_context
+    def try
+      yield
+    ensure
+      @ensure_context = self
+    end
+  end
+
+  yielder = EnsureYieldBreak.new
+  yielder.try do
+  end
+  assert_equal yielder, yielder.ensure_context
+end
+
+assert('ensure - context - yield and break') do
+  class EnsureYieldBreak
+    attr_reader :ensure_context
+    def try
+      yield
+    ensure
+      @ensure_context = self
+    end
+  end
+
+  yielder = EnsureYieldBreak.new
+  yielder.try do
+    break
+  end
+  assert_equal yielder, yielder.ensure_context
+end
+
+assert('ensure - context - yield and return') do
+  class EnsureYieldBreak
+    attr_reader :ensure_context
+    def try
+      yield
+    ensure
+      @ensure_context = self
+    end
+  end
+
+  yielder = EnsureYieldBreak.new
+  lambda do
+    yielder.try do
+      return
+    end
+  end.call
+  assert_equal yielder, yielder.ensure_context
+end