[PRISM] CompileEnsureNode

3 files changed, 124 insertions, 45 deletions

diff --git a/compile.c b/compile.c
index 9d09bef494..b69fbbe3be 100644
--- a/compile.c
+++ b/compile.c
@@ -120,7 +120,7 @@ struct ensure_range {
 };
 
 struct iseq_compile_data_ensure_node_stack {
-    const NODE *ensure_node;
+    const void *ensure_node;
     struct iseq_compile_data_ensure_node_stack *prev;
     struct ensure_range *erange;
 };
@@ -5750,7 +5750,7 @@ make_name_for_block(const rb_iseq_t *orig_iseq)
 static void
 push_ensure_entry(rb_iseq_t *iseq,
                   struct iseq_compile_data_ensure_node_stack *enl,
-                  struct ensure_range *er, const NODE *const node)
+                  struct ensure_range *er, const void *const node)
 {
     enl->ensure_node = node;
     enl->prev = ISEQ_COMPILE_DATA(iseq)->ensure_node_stack;	/* prev */
diff --git a/prism_compile.c b/prism_compile.c
index ed0b5b0709..e8be7e6bb7 100644
--- a/prism_compile.c
+++ b/prism_compile.c
@@ -12,7 +12,10 @@
 #define NEW_CHILD_ISEQ(node, name, type, line_no) \
     pm_new_child_iseq(iseq, (node), parser, rb_fstring(name), iseq, (type), (line_no))
 
-#define PM_COMPILE(_ret, node) \
+#define PM_COMPILE(node) \
+    pm_compile_node(iseq, (node), ret, src, popped, scope_node)
+
+#define PM_COMPILE_INTO_ANCHOR(_ret, node) \
     pm_compile_node(iseq, (node), _ret, src, popped, scope_node)
 
 #define PM_COMPILE_POPPED(node) \
@@ -436,7 +439,7 @@ pm_compile_flip_flop(pm_flip_flop_node_t *flip_flop_node, LABEL *else_label, LAB
     ADD_INSNL(ret, &dummy_line_node, branchif, lend);
 
     if (flip_flop_node->left) {
-        PM_COMPILE(ret, flip_flop_node->left);
+        PM_COMPILE(flip_flop_node->left);
     }
     else {
         PM_PUTNIL;
@@ -451,7 +454,7 @@ pm_compile_flip_flop(pm_flip_flop_node_t *flip_flop_node, LABEL *else_label, LAB
 
     ADD_LABEL(ret, lend);
     if (flip_flop_node->right) {
-        PM_COMPILE(ret, flip_flop_node->right);
+        PM_COMPILE(flip_flop_node->right);
     }
     else {
         PM_PUTNIL;
@@ -740,7 +743,7 @@ pm_compile_class_path(LINK_ANCHOR *const ret, rb_iseq_t *iseq, const pm_node_t *
         pm_node_t *parent = ((pm_constant_path_node_t *)constant_path_node)->parent;
         if (parent) {
             /* Bar::Foo */
-            PM_COMPILE(ret, parent);
+            PM_COMPILE(parent);
             return VM_DEFINECLASS_FLAG_SCOPED;
         }
         else {
@@ -1074,6 +1077,19 @@ pm_scope_node_init(const pm_node_t *node, pm_scope_node_t *scope, pm_scope_node_
             scope->locals = cast->locals;
             break;
         }
+        case PM_ENSURE_NODE: {
+            pm_ensure_node_t *cast = (pm_ensure_node_t *)node;
+
+            // HAXXX
+            pm_constant_id_t err_info = 3863;
+            pm_constant_id_list_t locals;
+            pm_constant_id_list_init(&locals);
+            pm_constant_id_list_append(&locals, err_info);
+
+            scope->body = (pm_node_t *)cast->statements;
+            scope->locals = locals;
+            break;
+        }
         case PM_FOR_NODE: {
             pm_for_node_t *cast = (pm_for_node_t *)node;
             scope->body = (pm_node_t *)cast->statements;
@@ -1195,7 +1211,7 @@ pm_compile_defined_expr0(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *co
           }
           pm_compile_defined_expr0(iseq, constant_path_node->parent, ret, src, popped, scope_node, dummy_line_node, lineno, true, lfinish);
           ADD_INSNL(ret, &dummy_line_node, branchunless, lfinish[1]);
-          PM_COMPILE(ret, constant_path_node->parent);
+          PM_COMPILE(constant_path_node->parent);
         } else {
           ADD_INSN1(ret, &dummy_line_node, putobject, rb_cObject);
         }
@@ -1317,8 +1333,8 @@ pm_setup_args(pm_arguments_node_t *arguments_node, int *flags, struct rb_callinf
                             case PM_ASSOC_NODE: {
                                 pm_assoc_node_t *assoc = (pm_assoc_node_t *)cur_node;
 
-                                PM_COMPILE(ret, assoc->key);
-                                PM_COMPILE(ret, assoc->value);
+                                PM_COMPILE(assoc->key);
+                                PM_COMPILE(assoc->value);
                                 cur_hash_size++;
 
                                 // If we're at the last keyword arg, or the last assoc node of this "set",
@@ -1349,7 +1365,7 @@ pm_setup_args(pm_arguments_node_t *arguments_node, int *flags, struct rb_callinf
                                 }
 
                                 pm_assoc_splat_node_t *assoc_splat = (pm_assoc_splat_node_t *)cur_node;
-                                PM_COMPILE(ret, assoc_splat->value);
+                                PM_COMPILE(assoc_splat->value);
 
                                 *flags |= VM_CALL_KW_SPLAT | VM_CALL_KW_SPLAT_MUT;
 
@@ -1379,7 +1395,7 @@ pm_setup_args(pm_arguments_node_t *arguments_node, int *flags, struct rb_callinf
                       for (size_t i = 0; i < len; i++) {
                           pm_assoc_node_t *assoc = (pm_assoc_node_t *)keyword_arg->elements.nodes[i];
                           (*kw_arg)->keywords[i] = pm_static_literal_value(assoc->key, scope_node, parser);
-                          PM_COMPILE(ret, assoc->value);
+                          PM_COMPILE(assoc->value);
                       }
                   }
                   break;
@@ -1388,7 +1404,7 @@ pm_setup_args(pm_arguments_node_t *arguments_node, int *flags, struct rb_callinf
                   *flags |= VM_CALL_ARGS_SPLAT;
                   pm_splat_node_t *splat_node = (pm_splat_node_t *)argument;
                   if (splat_node->expression) {
-                      PM_COMPILE(ret, splat_node->expression);
+                      PM_COMPILE(splat_node->expression);
                   }
 
                   if (!popped) {
@@ -1482,7 +1498,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         ADD_INSNL(ret, &dummy_line_node, branchunless, end_label);
 
         PM_POP_UNLESS_POPPED;
-        PM_COMPILE(ret, and_node->right);
+        PM_COMPILE(and_node->right);
         ADD_LABEL(ret, end_label);
         return;
       }
@@ -1490,7 +1506,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         pm_arguments_node_t *arguments_node = (pm_arguments_node_t *) node;
         pm_node_list_t node_list = arguments_node->arguments;
         for (size_t index = 0; index < node_list.size; index++) {
-            PM_COMPILE(ret, node_list.nodes[index]);
+            PM_COMPILE(node_list.nodes[index]);
         }
         return;
       }
@@ -1519,7 +1535,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
             pm_node_list_t *elements = &cast->elements;
 
             for (size_t index = 0; index < elements->size; index++) {
-                PM_COMPILE(ret, elements->nodes[index]);
+                PM_COMPILE(elements->nodes[index]);
             }
 
             if (!popped) {
@@ -1531,16 +1547,16 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
       }
       case PM_ASSOC_NODE: {
         pm_assoc_node_t *assoc_node = (pm_assoc_node_t *) node;
-        PM_COMPILE(ret, assoc_node->key);
+        PM_COMPILE(assoc_node->key);
         if (assoc_node->value) {
-            PM_COMPILE(ret, assoc_node->value);
+            PM_COMPILE(assoc_node->value);
         }
         return;
       }
       case PM_ASSOC_SPLAT_NODE: {
         pm_assoc_splat_node_t *assoc_splat_node = (pm_assoc_splat_node_t *)node;
 
-        PM_COMPILE(ret, assoc_splat_node->value);
+        PM_COMPILE(assoc_splat_node->value);
         return;
       }
       case PM_BACK_REFERENCE_READ_NODE: {
@@ -1555,18 +1571,61 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
       }
       case PM_BEGIN_NODE: {
         pm_begin_node_t *begin_node = (pm_begin_node_t *) node;
+        rb_iseq_t *child_iseq;
+        LABEL *lstart = NEW_LABEL(lineno);
+        LABEL *lend = NEW_LABEL(lineno);
+
+        ADD_LABEL(ret, lstart);
         if (begin_node->statements) {
-            PM_COMPILE(ret, (pm_node_t *)begin_node->statements);
+            PM_COMPILE((pm_node_t *)begin_node->statements);
         }
         else {
             PM_PUTNIL_UNLESS_POPPED;
         }
+        ADD_LABEL(ret, lend);
+
+        if (begin_node->ensure_clause) {
+            DECL_ANCHOR(ensr);
+
+            iseq_set_exception_local_table(iseq);
+            pm_scope_node_t next_scope_node;
+            pm_scope_node_init((pm_node_t *)begin_node->ensure_clause, &next_scope_node, scope_node, parser);
+
+            child_iseq = NEW_CHILD_ISEQ(next_scope_node,
+                    rb_str_new2("ensure in"),
+                    ISEQ_TYPE_ENSURE, lineno);
+
+            LABEL *lcont = NEW_LABEL(lineno);
+            struct ensure_range er;
+            struct iseq_compile_data_ensure_node_stack enl;
+            struct ensure_range *erange;
+
+            INIT_ANCHOR(ensr);
+            PM_COMPILE_INTO_ANCHOR(ensr, (pm_node_t *)&next_scope_node);
+
+            er.begin = lstart;
+            er.end = lend;
+            er.next = 0;
+            push_ensure_entry(iseq, &enl, &er, (void *)&next_scope_node);
+
+            ISEQ_COMPILE_DATA(iseq)->current_block = child_iseq;
+            ADD_SEQ(ret, ensr);
+            ADD_LABEL(ret, lcont);
+            erange = ISEQ_COMPILE_DATA(iseq)->ensure_node_stack->erange;
+            if (lstart->link.next != &lend->link) {
+                while (erange) {
+                    ADD_CATCH_ENTRY(CATCH_TYPE_ENSURE, erange->begin, erange->end, child_iseq, lcont);
+                    erange = erange->next;
+                }
+            }
+            PM_POP_UNLESS_POPPED;
+        }
         return;
       }
       case PM_BLOCK_ARGUMENT_NODE: {
         pm_block_argument_node_t *block_argument_node = (pm_block_argument_node_t *) node;
         if (block_argument_node->expression) {
-            PM_COMPILE(ret, block_argument_node->expression);
+            PM_COMPILE(block_argument_node->expression);
         }
         return;
       }
@@ -1733,7 +1792,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
 
             pm_while_node_t *condition_node = (pm_while_node_t *)conditions.nodes[i];
             if (condition_node->statements) {
-                PM_COMPILE(ret, (pm_node_t *)condition_node->statements);
+                PM_COMPILE((pm_node_t *)condition_node->statements);
             }
             else {
                 PM_PUTNIL_UNLESS_POPPED;
@@ -2005,7 +2064,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         pm_constant_path_target_node_t *cast = (pm_constant_path_target_node_t *)node;
 
         if (cast->parent) {
-            PM_COMPILE(ret, cast->parent);
+            PM_COMPILE(cast->parent);
         }
 
         return;
@@ -2164,7 +2223,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         pm_embedded_statements_node_t *embedded_statements_node = (pm_embedded_statements_node_t *)node;
 
         if (embedded_statements_node->statements) {
-            PM_COMPILE(ret, (pm_node_t *) (embedded_statements_node->statements));
+            PM_COMPILE((pm_node_t *) (embedded_statements_node->statements));
         }
         else {
             PM_PUTNIL;
@@ -2176,7 +2235,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
       }
       case PM_EMBEDDED_VARIABLE_NODE: {
         pm_embedded_variable_node_t *embedded_node = (pm_embedded_variable_node_t *)node;
-        PM_COMPILE(ret, embedded_node->variable);
+        PM_COMPILE(embedded_node->variable);
         return;
       }
       case PM_FALSE_NODE:
@@ -2187,7 +2246,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
       case PM_ELSE_NODE: {
           pm_else_node_t *cast = (pm_else_node_t *)node;
           if (cast->statements) {
-              PM_COMPILE(ret, (pm_node_t *)cast->statements);
+              PM_COMPILE((pm_node_t *)cast->statements);
           }
           else {
               PM_PUTNIL;
@@ -2423,7 +2482,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
 
                         ADD_INSN1(ret, &dummy_line_node, putspecialobject, INT2FIX(VM_SPECIAL_OBJECT_VMCORE));
                         PM_SWAP;
-                        PM_COMPILE(ret, elements->nodes[index]);
+                        PM_COMPILE(elements->nodes[index]);
 
                         allocated_hashes++;
                         elements_of_cur_type = 0;
@@ -2431,11 +2490,11 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
                     }
                     else {
                         elements_of_cur_type++;
-                        PM_COMPILE(ret, elements->nodes[index]);
+                        PM_COMPILE(elements->nodes[index]);
                     }
                 }
                 else {
-                    PM_COMPILE(ret, elements->nodes[index]);
+                    PM_COMPILE(elements->nodes[index]);
                 }
             }
 
@@ -2482,7 +2541,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         // of the missing value. To compile these nodes, we simply compile the
         // value that is implied, which is helpfully supplied by the parser.
         pm_implicit_node_t *cast = (pm_implicit_node_t *)node;
-        PM_COMPILE(ret, cast->value);
+        PM_COMPILE(cast->value);
         return;
       }
       case PM_INSTANCE_VARIABLE_AND_WRITE_NODE: {
@@ -2654,7 +2713,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         pm_node_list_t elements = keyword_hash_node->elements;
 
         for (size_t index = 0; index < elements.size; index++) {
-            PM_COMPILE(ret, elements.nodes[index]);
+            PM_COMPILE(elements.nodes[index]);
         }
 
         if (!popped) {
@@ -2970,7 +3029,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
                     ADD_INSN1(ret, &dummy_line_node, topn, INT2FIX(pushed));
                     ADD_INSN1(ret, &dummy_line_node, setconstant, ID2SYM(name));
                 } else {
-                    PM_COMPILE(ret, lefts->nodes[index]);
+                    PM_COMPILE(lefts->nodes[index]);
                 }
             }
         }
@@ -2988,14 +3047,14 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
             pm_splat_node_t *rest_splat = ((pm_splat_node_t *)multi_write_node->rest);
             if (rest_splat->expression) {
                 ADD_INSN2(ret, &dummy_line_node, expandarray, INT2FIX(0), INT2FIX(1));
-                PM_COMPILE(ret, rest_splat->expression);
+                PM_COMPILE(rest_splat->expression);
             }
         }
 
         if (rights->size) {
             ADD_INSN2(ret, &dummy_line_node, expandarray, INT2FIX(rights->size), INT2FIX(2));
             for (size_t index = 0; index < rights->size; index++) {
-                PM_COMPILE(ret, rights->nodes[index]);
+                PM_COMPILE(rights->nodes[index]);
             }
         }
 
@@ -3039,7 +3098,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         ADD_INSNL(ret, &dummy_line_node, branchif, end_label);
 
         PM_POP_UNLESS_POPPED;
-        PM_COMPILE(ret, or_node->right);
+        PM_COMPILE(or_node->right);
         ADD_LABEL(ret, end_label);
 
         return;
@@ -3065,7 +3124,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         if (parentheses_node->body == NULL) {
             PM_PUTNIL_UNLESS_POPPED;
         } else {
-            PM_COMPILE(ret, parentheses_node->body);
+            PM_COMPILE(parentheses_node->body);
         }
 
         return;
@@ -3144,13 +3203,13 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
             if (range_node->left == NULL) {
                 PM_PUTNIL;
             } else {
-                PM_COMPILE(ret, range_node->left);
+                PM_COMPILE(range_node->left);
             }
 
             if (range_node->right == NULL) {
                 PM_PUTNIL;
             } else {
-                PM_COMPILE(ret, range_node->right);
+                PM_COMPILE(range_node->right);
             }
 
             if (!popped) {
@@ -3183,7 +3242,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         pm_arguments_node_t *arguments = ((pm_return_node_t *)node)->arguments;
 
         if (arguments) {
-            PM_COMPILE(ret, (pm_node_t *)arguments);
+            PM_COMPILE((pm_node_t *)arguments);
         }
         else {
             PM_PUTNIL;
@@ -3302,7 +3361,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
                     pm_for_node_t *for_node = (pm_for_node_t *)scope_node->ast_node;
 
                     ADD_GETLOCAL(ret, &dummy_line_node, 1, 0);
-                    PM_COMPILE(ret, for_node->index);
+                    PM_COMPILE(for_node->index);
                     ADD_INSN(ret, &dummy_line_node, nop);
                   }
                   default: {
@@ -3323,9 +3382,22 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
             ADD_CATCH_ENTRY(CATCH_TYPE_NEXT, start, end, NULL, end);
             break;
         }
+        case ISEQ_TYPE_ENSURE: {
+            iseq_set_exception_local_table(iseq);
+            PM_COMPILE((pm_node_t *)scope_node->body);
+            PM_POP;
+
+            // this loses precision, but in this case that doesn't matter...
+            pm_constant_id_t id_err_info = (pm_constant_id_t)ISEQ_BODY(iseq)->local_table[0];
+            int local_index = pm_lookup_local_index_with_depth(iseq, scope_node, id_err_info, 0);
+            ADD_GETLOCAL(ret, &dummy_line_node, local_index, 0);
+            ADD_INSN1(ret, &dummy_line_node, throw, INT2FIX(0));
+
+            return;
+        }
         default:
             if (scope_node->body) {
-                PM_COMPILE(ret, (pm_node_t *)scope_node->body);
+                PM_COMPILE((pm_node_t *)scope_node->body);
             }
             else {
                 PM_PUTNIL;
@@ -3334,7 +3406,9 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
 
         free(index_lookup_table);
 
-        ADD_INSN(ret, &dummy_line_node, leave);
+        if (!scope_node->previous) {
+            ADD_INSN(ret, &dummy_line_node, leave);
+        }
         return;
       }
       case PM_SELF_NODE:
@@ -3396,7 +3470,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
       case PM_SPLAT_NODE: {
         pm_splat_node_t *splat_node = (pm_splat_node_t *)node;
         if (splat_node->expression) {
-            PM_COMPILE(ret, splat_node->expression);
+            PM_COMPILE(splat_node->expression);
         }
 
         if (!popped) {
@@ -3411,7 +3485,7 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
             for (size_t index = 0; index < node_list.size - 1; index++) {
                 PM_COMPILE_POPPED(node_list.nodes[index]);
             }
-            PM_COMPILE(ret, node_list.nodes[node_list.size - 1]);
+            PM_COMPILE(node_list.nodes[node_list.size - 1]);
         }
         else {
             PM_PUTNIL;
@@ -3420,8 +3494,8 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
       }
       case PM_STRING_CONCAT_NODE: {
         pm_string_concat_node_t *str_concat_node = (pm_string_concat_node_t *)node;
-        PM_COMPILE(ret, str_concat_node->left);
-        PM_COMPILE(ret, str_concat_node->right);
+        PM_COMPILE(str_concat_node->left);
+        PM_COMPILE(str_concat_node->right);
         if (!popped) {
             ADD_INSN1(ret, &dummy_line_node, concatstrings, INT2FIX(2));
         }
diff --git a/test/ruby/test_compile_prism.rb b/test/ruby/test_compile_prism.rb
index 911b126e73..f846009c05 100644
--- a/test/ruby/test_compile_prism.rb
+++ b/test/ruby/test_compile_prism.rb
@@ -587,6 +587,11 @@ module Prism
       assert_prism_eval("while true; break 1; end")
     end
 
+    def test_EnsureNode
+      assert_prism_eval("begin; 1; ensure; 2; end")
+      assert_prism_eval("begin; 1; begin; 3; ensure; 4; end; ensure; 2; end")
+    end
+
     def test_NextNode
       # TODO:
       # assert_prism_eval("2.times do |i|; next if i == 1; end")