[PRISM] Fix up ensure+loop+break

1 files changed, 8 insertions, 1 deletions

diff --git a/prism_compile.c b/prism_compile.c
index 3d3e2f8549..e423e189b8 100644
--- a/prism_compile.c
+++ b/prism_compile.c
@@ -1114,7 +1114,6 @@ pm_compile_loop(rb_iseq_t *iseq, const pm_node_location_t *node_location, pm_nod
     LABEL *prev_end_label = ISEQ_COMPILE_DATA(iseq)->end_label;
     LABEL *prev_redo_label = ISEQ_COMPILE_DATA(iseq)->redo_label;
 
-    // TODO: Deal with ensures in here
     LABEL *next_label = ISEQ_COMPILE_DATA(iseq)->start_label = NEW_LABEL(location.line); /* next  */
     LABEL *redo_label = ISEQ_COMPILE_DATA(iseq)->redo_label = NEW_LABEL(location.line);  /* redo  */
     LABEL *break_label = ISEQ_COMPILE_DATA(iseq)->end_label = NEW_LABEL(location.line);  /* break */
@@ -1124,6 +1123,12 @@ pm_compile_loop(rb_iseq_t *iseq, const pm_node_location_t *node_location, pm_nod
     LABEL *next_catch_label = NEW_LABEL(location.line);
     LABEL *tmp_label = NULL;
 
+    // We're pushing onto the ensure stack because breaks need to break out of
+    // this loop and not break into the ensure statements within the same
+    // lexical scope.
+    struct iseq_compile_data_ensure_node_stack enl;
+    push_ensure_entry(iseq, &enl, NULL, NULL);
+
     // begin; end while true
     if (flags & PM_LOOP_FLAGS_BEGIN_MODIFIER) {
         tmp_label = NEW_LABEL(location.line);
@@ -1176,6 +1181,8 @@ pm_compile_loop(rb_iseq_t *iseq, const pm_node_location_t *node_location, pm_nod
     ISEQ_COMPILE_DATA(iseq)->start_label = prev_start_label;
     ISEQ_COMPILE_DATA(iseq)->end_label = prev_end_label;
     ISEQ_COMPILE_DATA(iseq)->redo_label = prev_redo_label;
+    ISEQ_COMPILE_DATA(iseq)->ensure_node_stack = ISEQ_COMPILE_DATA(iseq)->ensure_node_stack->prev;
+
     return;
 }