[PRISM] Fixed redo node

ruby · Dec 4, 2023 · 81a7008 · 81a7008
1 parent 7d371ca
commit 81a7008
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 13 deletions.
diff --git a/prism_compile.c b/prism_compile.c
@@ -3722,7 +3722,60 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         return;
       }
       case PM_REDO_NODE: {
-        ADD_INSNL(ret, &dummy_line_node, jump, ISEQ_COMPILE_DATA(iseq)->redo_label);
+        if (ISEQ_COMPILE_DATA(iseq)->redo_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)->redo_label);
+
+            add_ensure_iseq(ret, iseq, 0);
+            ADD_INSNL(ret, &dummy_line_node, jump, ISEQ_COMPILE_DATA(iseq)->redo_label);
+            ADD_ADJUST_RESTORE(ret, splabel);
+            PM_PUTNIL_UNLESS_POPPED;
+        }
+        else if (ISEQ_BODY(iseq)->type != ISEQ_TYPE_EVAL && ISEQ_COMPILE_DATA(iseq)->start_label && 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)->start_label);
+            ADD_INSNL(ret, &dummy_line_node, jump, ISEQ_COMPILE_DATA(iseq)->start_label);
+            ADD_ADJUST_RESTORE(ret, splabel);
+
+            PM_PUTNIL_UNLESS_POPPED;
+        }
+        else {
+            const rb_iseq_t *ip = iseq;
+
+            while (ip) {
+              if (!ISEQ_COMPILE_DATA(ip)) {
+                  ip = 0;
+                  break;
+              }
+
+              if (ISEQ_COMPILE_DATA(ip)->redo_label != 0) {
+                  break;
+              }
+              else if (ISEQ_BODY(ip)->type == ISEQ_TYPE_BLOCK) {
+                  break;
+              }
+              else if (ISEQ_BODY(ip)->type == ISEQ_TYPE_EVAL) {
+                  rb_bug("Invalid redo\n");
+              }
+
+              ip = ISEQ_BODY(ip)->parent_iseq;
+          }
+          if (ip != 0) {
+              PM_PUTNIL;
+              ADD_INSN1(ret, &dummy_line_node, throw, INT2FIX(VM_THROW_NO_ESCAPE_FLAG | TAG_REDO));
+
+              PM_POP_IF_POPPED;
+          }
+          else {
+              rb_bug("Invalid redo\n");
+          }
+        }
         return;
       }
       case PM_REGULAR_EXPRESSION_NODE: {

diff --git a/test/ruby/test_compile_prism.rb b/test/ruby/test_compile_prism.rb
@@ -815,18 +815,33 @@ def test_NextNode
     end
 
     def test_RedoNode
-      # TODO:
-      # assert_prism_eval(<<-CODE
-      # counter = 0
-
-      # 5.times do |i|
-      #   counter += 1
-      #   if i == 2 && counter < 3
-      #     redo
-      #   end
-      # end
-      # CODE
-      # )
+      assert_prism_eval(<<-CODE)
+        counter = 0
+
+        5.times do |i|
+          counter += 1
+          if i == 2 && counter < 3
+            redo
+          end
+        end
+      CODE
+
+      assert_prism_eval(<<-CODE)
+        for i in 1..5
+          if i == 3
+            i = 0
+            redo
+          end
+        end
+      CODE
+
+      assert_prism_eval(<<-CODE)
+        i = 0
+        begin
+          i += 1
+          redo if i == 3
+        end while i < 5
+      CODE
     end
 
     def test_RescueNode