From d6ec1aa0a79235e75cbb92141590199a322e7bf1 Mon Sep 17 00:00:00 2001
From: Matt Valentine-House <matt@eightbitraptor.com>
Date: Fri, 1 Dec 2023 14:37:10 +0000
Subject: [PATCH] [PRISM] Compile Rescue Modifier nodes

---
 prism_compile.c                 | 52 ++++++++++++++++++++++++++++++++-
 test/ruby/test_compile_prism.rb |  9 +++++-
 2 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/prism_compile.c b/prism_compile.c
index 61f9b9f5499aac..841dae979681e8 100644
--- a/prism_compile.c
+++ b/prism_compile.c
@@ -1419,6 +1419,12 @@ pm_scope_node_init(const pm_node_t *node, pm_scope_node_t *scope, pm_scope_node_
             scope->local_depth_offset += 1;
             break;
         }
+        case PM_RESCUE_MODIFIER_NODE: {
+            pm_rescue_modifier_node_t *cast = (pm_rescue_modifier_node_t *)node;
+            scope->body = (pm_node_t *)cast->rescue_expression;
+            scope->local_depth_offset += 1;
+            break;
+        }
         case PM_SINGLETON_CLASS_NODE: {
             pm_singleton_class_node_t *cast = (pm_singleton_class_node_t *) node;
             scope->body = cast->body;
@@ -3842,6 +3848,34 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
 
         return;
       }
+      case PM_RESCUE_MODIFIER_NODE: {
+        pm_scope_node_t rescue_scope_node;
+        pm_rescue_modifier_node_t *rescue_node = (pm_rescue_modifier_node_t *)node;
+        pm_scope_node_init((pm_node_t *)rescue_node, &rescue_scope_node, scope_node, parser);
+
+        rb_iseq_t *rescue_iseq = NEW_CHILD_ISEQ(rescue_scope_node,
+                                                rb_str_concat(rb_str_new2("rescue in"),
+                                                              ISEQ_BODY(iseq)->location.label),
+                                                ISEQ_TYPE_RESCUE, 1);
+
+        LABEL *lstart = NEW_LABEL(lineno);
+        LABEL *lend = NEW_LABEL(lineno);
+        LABEL *lcont = NEW_LABEL(lineno);
+
+        lstart->rescued = LABEL_RESCUE_BEG;
+        lend->rescued = LABEL_RESCUE_END;
+        ADD_LABEL(ret, lstart);
+        PM_COMPILE_NOT_POPPED((pm_node_t *)rescue_node->expression);
+        ADD_LABEL(ret, lend);
+        ADD_INSN(ret, &dummy_line_node, nop);
+        ADD_LABEL(ret, lcont);
+
+        PM_POP_IF_POPPED;
+
+        ADD_CATCH_ENTRY(CATCH_TYPE_RESCUE, lstart, lend, rescue_iseq, lcont);
+        ADD_CATCH_ENTRY(CATCH_TYPE_RETRY, lend, lcont, NULL, lstart);
+        return;
+      }
       case PM_RETURN_NODE: {
         pm_arguments_node_t *arguments = ((pm_return_node_t *)node)->arguments;
 
@@ -4211,7 +4245,23 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         }
         case ISEQ_TYPE_RESCUE: {
             iseq_set_exception_local_table(iseq);
-            PM_COMPILE((pm_node_t *)scope_node->ast_node);
+            if (PM_NODE_TYPE_P(scope_node->ast_node, PM_RESCUE_MODIFIER_NODE)) {
+                LABEL *lab = NEW_LABEL(lineno);
+                LABEL *rescue_end = NEW_LABEL(lineno);
+                ADD_GETLOCAL(ret, &dummy_line_node, LVAR_ERRINFO, 0);
+                ADD_INSN1(ret, &dummy_line_node, putobject, rb_eStandardError);
+                ADD_INSN1(ret, &dummy_line_node, checkmatch, INT2FIX(VM_CHECKMATCH_TYPE_RESCUE));
+                ADD_INSN1(ret, &dummy_line_node, branchif, lab);
+                ADD_INSN1(ret, &dummy_line_node, jump, rescue_end);
+                ADD_LABEL(ret, lab);
+                PM_COMPILE((pm_node_t *)scope_node->body);
+                ADD_INSN(ret, &dummy_line_node, leave);
+                ADD_LABEL(ret, rescue_end);
+                ADD_GETLOCAL(ret, &dummy_line_node, LVAR_ERRINFO, 0);
+            }
+            else {
+                PM_COMPILE((pm_node_t *)scope_node->ast_node);
+            }
             ADD_INSN1(ret, &dummy_line_node, throw, INT2FIX(0));
 
             return;
diff --git a/test/ruby/test_compile_prism.rb b/test/ruby/test_compile_prism.rb
index 905e29b0d22f08..fb377813fb8f99 100644
--- a/test/ruby/test_compile_prism.rb
+++ b/test/ruby/test_compile_prism.rb
@@ -914,6 +914,14 @@ def test_RescueNode
       assert_prism_eval("begin; rescue; end")
     end
 
+    def test_RescueModiferNode
+      assert_prism_eval("1.nil? rescue false")
+      assert_prism_eval("1.nil? rescue 1")
+      assert_prism_eval("raise 'bang' rescue nil")
+      assert_prism_eval("raise 'bang' rescue a = 1; a.nil?")
+      assert_prism_eval("a = 0 rescue (a += 1 && retry if a <= 1)")
+    end
+
     def test_RetryNode
       assert_prism_eval(<<~CODE)
         a = 1
@@ -945,7 +953,6 @@ def test_RetryNode
       CODE
     end
 
-
     def test_ReturnNode
       assert_prism_eval("def return_node; return 1; end")
     end