[PRISM] Fix compilation for NextNode

This code was almost enitrely the same as the existing compiler's
code for its NextNode.

2 files changed, 104 insertions, 8 deletions

diff --git a/prism_compile.c b/prism_compile.c
index 79e4e66f86..2c7b4f265f 100644
--- a/prism_compile.c
+++ b/prism_compile.c
@@ -3460,15 +3460,81 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
       }
       case PM_NEXT_NODE: {
         pm_next_node_t *next_node = (pm_next_node_t *) node;
-        if (next_node->arguments) {
-            PM_COMPILE_NOT_POPPED((pm_node_t *)next_node->arguments);
+
+        if (ISEQ_COMPILE_DATA(iseq)->redo_label != 0 && can_add_ensure_iseq(iseq)) {
+            LABEL *splabel = NEW_LABEL(0);
+
+            ADD_LABEL(ret, splabel);
+
+            add_ensure_iseq(ret, iseq, 0);
+
+            ADD_ADJUST(ret, &dummy_line_node, ISEQ_COMPILE_DATA(iseq)->redo_label);
+            ADD_INSNL(ret, &dummy_line_node, jump, ISEQ_COMPILE_DATA(iseq)->start_label);
+
+            ADD_ADJUST_RESTORE(ret, splabel);
+            PM_PUTNIL_UNLESS_POPPED;
         }
-        else {
-            PM_PUTNIL;
+        else if (ISEQ_COMPILE_DATA(iseq)->end_label && can_add_ensure_iseq(iseq)) {
+            LABEL *splabel = NEW_LABEL(0);
+
+            ADD_LABEL(ret, splabel);
+            ADD_ADJUST(ret, &dummy_line_node, ISEQ_COMPILE_DATA(iseq)->start_label);
+
+            if (next_node->arguments) {
+                PM_COMPILE((pm_node_t *)next_node->arguments);
+            }
+            else {
+                PM_PUTNIL;
+            }
+
+            add_ensure_iseq(ret, iseq, 0);
+            ADD_INSNL(ret, &dummy_line_node, jump, ISEQ_COMPILE_DATA(iseq)->end_label);
+            ADD_ADJUST_RESTORE(ret, splabel);
+            splabel->unremovable = FALSE;
+
+            PM_PUTNIL_UNLESS_POPPED;
         }
+        else {
+            const rb_iseq_t *ip = iseq;
 
-        PM_POP;
-        ADD_INSNL(ret, &dummy_line_node, jump, ISEQ_COMPILE_DATA(iseq)->start_label);
+            unsigned long throw_flag = 0;
+            while (ip) {
+                if (!ISEQ_COMPILE_DATA(ip)) {
+                    ip = 0;
+                    break;
+                }
+
+                throw_flag = VM_THROW_NO_ESCAPE_FLAG;
+                if (ISEQ_COMPILE_DATA(ip)->redo_label != 0) {
+                    /* while loop */
+                    break;
+                }
+                else if (ISEQ_BODY(ip)->type == ISEQ_TYPE_BLOCK) {
+                    break;
+                }
+                else if (ISEQ_BODY(ip)->type == ISEQ_TYPE_EVAL) {
+                    rb_raise(rb_eArgError, "Can't escape from eval with next");
+                    return;
+                }
+
+                ip = ISEQ_BODY(ip)->parent_iseq;
+            }
+            if (ip != 0) {
+                if (next_node->arguments) {
+                    PM_COMPILE((pm_node_t *)next_node->arguments);
+                }
+                else {
+                    PM_PUTNIL;
+                }
+                ADD_INSN1(ret, &dummy_line_node, throw, INT2FIX(throw_flag | TAG_NEXT));
+
+                PM_POP_IF_POPPED;
+            }
+            else {
+                rb_raise(rb_eArgError, "Invalid next");
+                return;
+            }
+        }
 
         return;
       }
diff --git a/test/ruby/test_compile_prism.rb b/test/ruby/test_compile_prism.rb
index 34abf982ae..e51ca190d5 100644
--- a/test/ruby/test_compile_prism.rb
+++ b/test/ruby/test_compile_prism.rb
@@ -760,8 +760,38 @@ module Prism
     end
 
     def test_NextNode
-      # TODO:
-      # assert_prism_eval("2.times do |i|; next if i == 1; end")
+      assert_prism_eval("2.times do |i|; next if i == 1; end")
+
+      assert_prism_eval(<<-CODE)
+        res = []
+        i = 0
+        while i < 5
+          i += 1
+          next if i == 3
+          res << i
+        end
+        res
+      CODE
+
+      assert_prism_eval(<<-CODE)
+        res = []
+        (1..5).each do |i|
+          next if i.even?
+          res << i
+        end
+        res
+      CODE
+
+      assert_prism_eval(<<-CODE)
+        res = []
+        i = 0
+        begin
+          i += 1
+          next if i == 3
+          res << i
+        end while i < 5
+        res
+      CODE
     end
 
     def test_RedoNode