[PRISM] Fix ReturnNodes

This code is almost exactly the same (fixed variable names) as
what exists already in compile.c

2 files changed, 58 insertions, 10 deletions

diff --git a/prism_compile.c b/prism_compile.c
index 7c41f5804b..3843cfb9ea 100644
--- a/prism_compile.c
+++ b/prism_compile.c
@@ -3902,19 +3902,59 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
       case PM_RETURN_NODE: {
         pm_arguments_node_t *arguments = ((pm_return_node_t *)node)->arguments;
 
-        if (arguments) {
-            PM_COMPILE((pm_node_t *)arguments);
-        }
-        else {
-            PM_PUTNIL;
-        }
+        if (iseq) {
+            enum rb_iseq_type type = ISEQ_BODY(iseq)->type;
+            const NODE *retval = RNODE_RETURN(node)->nd_stts;
+            LABEL *splabel = 0;
+
+            const rb_iseq_t *parent_iseq = iseq;
+            enum rb_iseq_type parent_type = ISEQ_BODY(parent_iseq)->type;
+            while (parent_type == ISEQ_TYPE_RESCUE || parent_type == ISEQ_TYPE_ENSURE) {
+                if (!(parent_iseq = ISEQ_BODY(parent_iseq)->parent_iseq)) break;
+                parent_type = ISEQ_BODY(parent_iseq)->type;
+            }
 
-        ADD_TRACE(ret, RUBY_EVENT_RETURN);
-        ADD_INSN(ret, &dummy_line_node, leave);
+            switch (parent_type) {
+              case ISEQ_TYPE_TOP:
+              case ISEQ_TYPE_MAIN:
+                if (retval) {
+                    rb_warn("argument of top-level return is ignored");
+                }
+                if (parent_iseq == iseq) {
+                    type = ISEQ_TYPE_METHOD;
+                }
+                break;
+              default:
+                break;
+            }
 
-        if (!popped) {
-            PM_PUTNIL;
+            if (type == ISEQ_TYPE_METHOD) {
+                splabel = NEW_LABEL(0);
+                ADD_LABEL(ret, splabel);
+                ADD_ADJUST(ret, &dummy_line_node, 0);
+            }
+
+            if (arguments) {
+                PM_COMPILE((pm_node_t *)arguments);
+            }
+            else {
+                PM_PUTNIL;
+            }
+
+            if (type == ISEQ_TYPE_METHOD && can_add_ensure_iseq(iseq)) {
+                add_ensure_iseq(ret, iseq, 1);
+                ADD_TRACE(ret, RUBY_EVENT_RETURN);
+                ADD_INSN(ret, &dummy_line_node, leave);
+                ADD_ADJUST_RESTORE(ret, splabel);
+
+                PM_PUTNIL_UNLESS_POPPED;
+            }
+            else {
+                ADD_INSN1(ret, &dummy_line_node, throw, INT2FIX(TAG_RETURN));
+                PM_POP_IF_POPPED;
+            }
         }
+
         return;
       }
       case PM_RETRY_NODE: {
diff --git a/test/ruby/test_compile_prism.rb b/test/ruby/test_compile_prism.rb
index 62edb329bc..a707aa5a9c 100644
--- a/test/ruby/test_compile_prism.rb
+++ b/test/ruby/test_compile_prism.rb
@@ -967,6 +967,14 @@ module Prism
 
     def test_ReturnNode
       assert_prism_eval("def return_node; return 1; end")
+      assert_prism_eval(<<-CODE)
+        def self.prism_test_return_node
+          [1].each do |e|
+            return true
+          end
+        end
+        prism_test_return_node
+      CODE
     end
 
     ############################################################################