From 342f10b9b3579378d0983312c1f52c4c23029283 Mon Sep 17 00:00:00 2001
From: nobu <nobu@b2dd03c8-39d4-4d8f-98ff-823fe69b080e>
Date: Wed, 21 Dec 2016 01:58:32 +0000
Subject: [PATCH] compile.c: toplevel return

* compile.c (iseq_compile_each): stop execution of the current source
  by toplevel return.  [ruby-core:36785] [Feature #4840]

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@57132 b2dd03c8-39d4-4d8f-98ff-823fe69b080e
---
 NEWS                     |  2 ++
 compile.c                | 32 ++++++++++++++++++++++++--------
 test/ruby/test_syntax.rb | 23 +++++++++++++++++++++++
 3 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/NEWS b/NEWS
index 5c26678b0e..1c0bc904e4 100644
--- a/NEWS
+++ b/NEWS
@@ -25,6 +25,8 @@ with all sufficient information, see the ChangeLog file or Redmine
 * Rescue modifier now applicable to method arguments.
   [Feature #12686]
 
+* Toplevel return is now allowed.  [Feature #4840]
+
 === Core classes updates (outstanding ones only)
 
 * Array
diff --git a/compile.c b/compile.c
index d437be3028..c8e8add49a 100644
--- a/compile.c
+++ b/compile.c
@@ -4635,12 +4635,18 @@ iseq_compile_each(rb_iseq_t *iseq, LINK_ANCHOR *const ret, NODE *node, int poppe
 	LABEL *lstart = NEW_LABEL(line);
 	LABEL *lend = NEW_LABEL(line);
 	LABEL *lcont = NEW_LABEL(line);
+	LINK_ELEMENT *last;
+	int last_leave = 0;
 	struct ensure_range er;
 	struct iseq_compile_data_ensure_node_stack enl;
 	struct ensure_range *erange;
 
 	INIT_ANCHOR(ensr);
 	COMPILE_POPPED(ensr, "ensure ensr", node->nd_ensr);
+	last = ensr->last;
+	last_leave = last && IS_INSN(last) && IS_INSN_ID(last, leave);
+	if (!popped && last_leave)
+	    popped = 1;
 
 	er.begin = lstart;
 	er.end = lend;
@@ -4657,12 +4663,15 @@ iseq_compile_each(rb_iseq_t *iseq, LINK_ANCHOR *const ret, NODE *node, int poppe
 	    ADD_SEQ(ret, ensr);
 	}
 	ADD_LABEL(ret, lcont);
+	if (last_leave) ADD_INSN(ret, line, pop);
 
 	erange = ISEQ_COMPILE_DATA(iseq)->ensure_node_stack->erange;
-	while (erange) {
-	    ADD_CATCH_ENTRY(CATCH_TYPE_ENSURE, erange->begin, erange->end,
-			    ensure, lcont);
-	    erange = erange->next;
+	if (lstart->link.next != &lend->link) {
+	    while (erange) {
+		ADD_CATCH_ENTRY(CATCH_TYPE_ENSURE, erange->begin, erange->end,
+				ensure, lcont);
+		erange = erange->next;
+	    }
 	}
 
 	ISEQ_COMPILE_DATA(iseq)->ensure_node_stack = enl.prev;
@@ -5463,13 +5472,20 @@ iseq_compile_each(rb_iseq_t *iseq, LINK_ANCHOR *const ret, NODE *node, int poppe
 	rb_iseq_t *is = iseq;
 
 	if (is) {
-	    if (is->body->type == ISEQ_TYPE_TOP) {
-		COMPILE_ERROR(ERROR_ARGS "Invalid return");
+	    enum iseq_type type = is->body->type;
+	    const rb_iseq_t *parent_iseq = is->body->parent_iseq;
+	    enum iseq_type parent_type = parent_iseq ? parent_iseq->body->type : type;
+
+	    if (type == ISEQ_TYPE_TOP || type == ISEQ_TYPE_MAIN ||
+		((type == ISEQ_TYPE_RESCUE || type == ISEQ_TYPE_ENSURE) &&
+		 (parent_type == ISEQ_TYPE_TOP || parent_type == ISEQ_TYPE_MAIN))) {
+		ADD_INSN(ret, line, putnil);
+		ADD_INSN(ret, line, leave);
 	    }
 	    else {
 		LABEL *splabel = 0;
 
-		if (is->body->type == ISEQ_TYPE_METHOD) {
+		if (type == ISEQ_TYPE_METHOD) {
 		    splabel = NEW_LABEL(0);
 		    ADD_LABEL(ret, splabel);
 		    ADD_ADJUST(ret, line, 0);
@@ -5477,7 +5493,7 @@ iseq_compile_each(rb_iseq_t *iseq, LINK_ANCHOR *const ret, NODE *node, int poppe
 
 		COMPILE(ret, "return nd_stts (return val)", node->nd_stts);
 
-		if (is->body->type == ISEQ_TYPE_METHOD) {
+		if (type == ISEQ_TYPE_METHOD) {
 		    add_ensure_iseq(ret, iseq, 1);
 		    ADD_TRACE(ret, line, RUBY_EVENT_RETURN);
 		    ADD_INSN(ret, line, leave);
diff --git a/test/ruby/test_syntax.rb b/test/ruby/test_syntax.rb
index f3db9bcd48..3e6212e27c 100644
--- a/test/ruby/test_syntax.rb
+++ b/test/ruby/test_syntax.rb
@@ -926,6 +926,29 @@ eom
     assert_equal(:ok, result)
   end
 
+  def test_return_toplevel
+    feature4840 = '[ruby-core:36785] [Feature #4840]'
+    code = "#{<<~"begin;"}\n#{<<~"end;"}"
+    begin;
+      return; raise
+      begin return; rescue SystemExit; exit false; end
+      begin return; ensure exit false; end
+      begin ensure return; end
+      begin raise; ensure; return; end
+      begin raise; rescue; return; end
+      return false; raise
+      return 1; raise
+    end;
+    all_assertions(feature4840) do |a|
+      code.each_line do |s|
+        s.chomp!
+        a.for(s) do
+          assert_ruby_status([], s, proc {RubyVM::InstructionSequence.compile(s).disasm})
+        end
+      end
+    end
+  end
+
   private
 
   def not_label(x) @result = x; @not_label ||= nil end