Run ensure when return exits a begin/ensure

lib/natalie/compiler/env_builder.rb

@@ -73,7 +73,7 @@ def process_end_if(_)
       end
 
       def process_try(i)
-        @env = i.env || { vars: {}, outer: @env, hoist: true, type: i.label }
+        @env = i.env || { vars: {}, outer: @env, hoist: true, type: i.label, has_ensure: i.has_ensure }
       end
       def process_end_try(_)
         @env = @env.fetch(:outer)

lib/natalie/compiler/instructions/try_instruction.rb

@@ -3,16 +3,18 @@
 module Natalie
   class Compiler
     class TryInstruction < BaseInstruction
-      attr_reader :serial
+      attr_reader :serial, :has_ensure
 
-      def initialize(discard_catch_result: false)
+      def initialize(discard_catch_result: false, has_ensure: false)
         @discard_catch_result = discard_catch_result
+        @has_ensure = has_ensure
         @serial = self.class.serial
       end
 
       def to_s
         s = 'try'
         s << ' (discard_catch_result)' if @discard_catch_result
+        s << ' (ensure)' if @has_ensure
         s
       end
 
@@ -50,13 +52,28 @@ def generate(transform)
           code << '} catch(ExceptionObject *exception) {'
 
           code << 'auto exception_was = env->exception()'
-          code << 'env->set_exception(exception)'
+          # The Defer ensures $! is restored even when the catch body itself
+          # throws (e.g. a `return` from inside rescue, which propagates via
+          # LocalJumpError so a surrounding ensure can run).
+          code << 'Defer restore_exception([&] { env->set_exception(exception_was); })'
+          if @has_ensure
+            # Ensure-only try: control-flow throws (LocalJumpError from
+            # `return`, `break`, etc.) must not clobber $! while the ensure
+            # body runs; only real Ruby exceptions become the new $!.
+            code << 'if (exception->local_jump_error_type() == LocalJumpErrorType::None) env->set_exception(exception)'
+          else
+            code << 'env->set_exception(exception)'
+          end
 
           transform.with_same_scope(catch_body) { |t| code << t.transform(@discard_catch_result ? nil : "#{result} =") }
 
-          code << 'env->set_exception(exception_was)'
           code << 'GlobalEnv::the()->set_rescued(true)'
 
+          # For an ensure-only try the catch body just runs the ensure clause;
+          # the original exception (whether a real Ruby exception or a
+          # control-flow LocalJumpError) must propagate up after that.
+          code << 'throw exception' if @has_ensure
+
           code << '}'
         end
 
@@ -79,11 +96,23 @@ def execute(vm)
           vm.rescued = false
         rescue => e
           vm.rescued = true
-          vm.global_variables[:$!] = e
+          exception_was = vm.global_variables[:$!]
+          # For an ensure-only try, control-flow throws (LocalJumpError from
+          # `return`, `break`, etc.) must not clobber $! while the ensure body
+          # runs; only real Ruby exceptions become the new $!.
+          vm.global_variables[:$!] = e unless @has_ensure && e.is_a?(LocalJumpError)
           vm.ip = catch_ip
           vm.run
           vm.ip = end_ip
-          vm.global_variables.delete(:$!)
+          if exception_was
+            vm.global_variables[:$!] = exception_was
+          else
+            vm.global_variables.delete(:$!)
+          end
+          # The catch body for an ensure-only try just runs the ensure clause;
+          # the original exception must propagate up after that. (For a rescue
+          # try, the catch body itself raises if no clause matched.)
+          raise e if @has_ensure
         end
       end
 

lib/natalie/compiler/pass1.rb

@@ -378,12 +378,15 @@ def transform_begin_node(node, used:)
         end
 
         if node.ensure_clause&.statements
-          raise_call = Prism.call_node(receiver: nil, name: :raise, location: node.ensure_clause.location)
-          instructions.unshift(TryInstruction.new(discard_catch_result: true))
+          # The TryInstruction with has_ensure: true re-throws the caught
+          # exception itself at the C++ level after the ensure body runs, so
+          # we don't include a synthetic `raise` here — that would re-raise
+          # $! (which is intentionally preserved, not set to the caught
+          # exception, so the ensure body sees the right value).
+          instructions.unshift(TryInstruction.new(discard_catch_result: true, has_ensure: true))
           instructions += [
             CatchInstruction.new,
             transform_expression(node.ensure_clause.statements, used: true),
-            transform_expression(raise_call, used: true),
             EndInstruction.new(:try),
             transform_expression(node.ensure_clause.statements, used: false),
           ]

lib/natalie/compiler/pass4.rb

@@ -23,10 +23,20 @@ def transform
       private
 
       def transform_return(instruction)
+        # Check whether there is an ensure-protected `try` between us and the
+        # enclosing method/block; if so, the return must throw so that the
+        # ensure's catch handler runs (a bare C++ `return` would skip ensure).
+        has_enclosing_ensure = false
+        cur = @env
+        until cur.nil? || %i[define_method define_block].include?(cur.fetch(:type, nil))
+          has_enclosing_ensure ||= cur[:has_ensure]
+          cur = cur[:outer]
+        end
+
         env = @env
         env = env[:outer] while env[:hoist]
 
-        unless env.fetch(:type) == :define_block
+        if env.fetch(:type) != :define_block && !has_enclosing_ensure
           # ReturnInstruction inside anything else is OK as-is.
           return
         end

spec/language/ensure_spec.rb

@@ -131,15 +131,11 @@
   end
 
   it "has an impact on the method's explicit return value" do
-    NATFIXME "it has an impact on the method's explicit return value", exception: SpecFailedException do
-      @obj.explicit_return_in_method_with_ensure.should == :ensure
-    end
+    @obj.explicit_return_in_method_with_ensure.should == :ensure
   end
 
   it "has an impact on the method's explicit return value from rescue if returns explicitly" do
-    NATFIXME "it has an impact on the method's explicit return value from rescue if returns explicitly", exception: SpecFailedException do
-      @obj.explicit_return_in_rescue_and_explicit_return_in_ensure.should == "returned in ensure"
-    end
+    @obj.explicit_return_in_rescue_and_explicit_return_in_ensure.should == "returned in ensure"
   end
 
   it "has no impact on the method's explicit return value from rescue if returns implicitly" do

spec/language/predefined_spec.rb

@@ -402,9 +402,7 @@ def foo
     rescue
       $!.should == outer
     end
-    NATFIXME 'should be set to the value of $! before the begin after a successful rescue within an ensure', exception: SpecFailedException do
-      $!.should == nil
-    end
+    $!.should == nil
   end
 
   it "should be set to the new exception after a throwing rescue" do
@@ -430,9 +428,7 @@ def foo
       rescue # do not make the exception fail the example
         $!.should == inner
       end
-      NATFIXME 'should be set to the new exception after a throwing rescue', exception: SpecFailedException do
-        $!.should == outer
-      end
+      $!.should == outer
     end
     $!.should == nil
   end
@@ -498,9 +494,7 @@ def foo
       rescue
         $!.should == e
       end
-      NATFIXME 'should not be cleared when an exception is not rescued', exception: SpecFailedException do
-        $!.should == nil
-      end
+      $!.should == nil
     end
 
     it "should not be cleared when an exception is rescued and rethrown" do
@@ -519,9 +513,7 @@ def foo
       rescue
         $!.should == e
       end
-      NATFIXME 'should not be cleared when an exception is rescued and rethrown', exception: SpecFailedException do
-        $!.should == nil
-      end
+      $!.should == nil
     end
   end
 
@@ -1092,7 +1084,7 @@ def obj.foo2; yield; end
   end
 
   it "can be changed via <<" do
-    NATFIXME 'Implement $:', exception: LoadError, message: "Cannot manipulate $: at runtime (spec/language/predefined_spec.rb#1096)" do
+    NATFIXME 'Implement $:', exception: LoadError, message: "Cannot manipulate $: at runtime (spec/language/predefined_spec.rb#1088)" do
       $: << "foo"
       $:.should include("foo")
     ensure

spec/language/return_spec.rb

@@ -90,9 +90,7 @@ def f()
         ScratchPad << :function
       end
       f().should == :begin
-      NATFIXME 'it executes ensure before returning', exception: SpecFailedException do
-        ScratchPad.recorded.should == [:begin, :ensure]
-      end
+      ScratchPad.recorded.should == [:begin, :ensure]
     end
 
     it "returns last value returned in ensure" do
@@ -108,10 +106,8 @@ def f()
         end
         ScratchPad << :function
       end
-      NATFIXME 'it returns last value returned in ensure', exception: SpecFailedException do
-        f().should == :ensure
-        ScratchPad.recorded.should == [:begin, :ensure]
-      end
+      f().should == :ensure
+      ScratchPad.recorded.should == [:begin, :ensure]
     end
 
     it "executes nested ensures before returning" do
@@ -133,9 +129,7 @@ def f()
         ScratchPad << :function
       end
       f().should == :inner_begin
-      NATFIXME 'it executes nested ensures before returning', exception: SpecFailedException do
-        ScratchPad.recorded.should == [:inner_begin, :inner_ensure, :outer_ensure]
-      end
+      ScratchPad.recorded.should == [:inner_begin, :inner_ensure, :outer_ensure]
     end
 
     it "returns last value returned in nested ensures" do
@@ -160,10 +154,8 @@ def f()
         end
         ScratchPad << :function
       end
-      NATFIXME 'it returns last value returned in nested ensures', exception: SpecFailedException do
-        f().should == :outer_ensure
-        ScratchPad.recorded.should == [:inner_begin, :inner_ensure, :outer_ensure]
-      end
+      f().should == :outer_ensure
+      ScratchPad.recorded.should == [:inner_begin, :inner_ensure, :outer_ensure]
     end
 
     it "executes the ensure clause when begin/ensure are inside a lambda" do
@@ -364,17 +356,15 @@ def f
       end
 
       it "fires ensure block before returning" do
-        NATFIXME 'it fires ensure block before returning', exception: SpecFailedException do
-          ruby_exe(<<-END_OF_CODE).should == "within ensure\n"
-            begin
-              return
-            ensure
-              puts "within ensure"
-            end
-
-            puts "after begin"
-          END_OF_CODE
-        end
+        ruby_exe(<<-END_OF_CODE).should == "within ensure\n"
+          begin
+            return
+          ensure
+            puts "within ensure"
+          end
+
+          puts "after begin"
+        END_OF_CODE
       end
 
       # NATFIXME: Compile time error, load with non-static value

src/main_template.cpp

@@ -27,16 +27,23 @@ extern "C" int EVAL(Env *env, Value *result_memory) {
 
     Value result = Value::nil();
     try {
-        // FIXME: top-level `return` in a Ruby script should probably be changed to `exit`.
         result = [&]() -> Value {
             /*NAT_EVAL_BODY*/
             return Value::nil();
         }();
         run_exit_handlers = false;
         run_at_exit_handlers(env);
     } catch (ExceptionObject *exception) {
-        handle_top_level_exception(env, exception, run_exit_handlers);
-        return 1;
+        // Top-level `return` reaches here as a LocalJumpError(:Return) when
+        // wrapped in a begin/ensure (the conversion is needed to run ensure).
+        // MRI exits the script gracefully and discards the value.
+        if (exception->local_jump_error_type() == LocalJumpErrorType::Return) {
+            run_exit_handlers = false;
+            run_at_exit_handlers(env);
+        } else {
+            handle_top_level_exception(env, exception, run_exit_handlers);
+            return 1;
+        }
     }
     memcpy(result_memory, &result, sizeof(Value));
     return 0;

test/natalie/return_test.rb

@@ -197,4 +197,69 @@ def define_it(name, &block)
     define_it(:buz, &proc { 1.tap { return 13 } })
     buz.should == 13
   end
+
+  describe 'inside a begin/ensure' do
+    it 'runs the ensure block before returning from the method' do
+      log = []
+      m = ->(_) {} # silence unused
+      _ = m
+      def f1(log)
+        begin
+          log << :body
+          return :body_value
+        ensure
+          log << :ensure
+        end
+      end
+      f1(log).should == :body_value
+      log.should == %i[body ensure]
+    end
+
+    it 'runs ensure before returning from a rescue clause' do
+      log = []
+      def f2(log)
+        begin
+          raise 'x'
+        rescue StandardError
+          log << :rescue
+          return :rescue_value
+        ensure
+          log << :ensure
+        end
+      end
+      f2(log).should == :rescue_value
+      log.should == %i[rescue ensure]
+    end
+
+    it 'lets a return in the ensure clause override the return value' do
+      def f3
+        begin
+          return 1
+        ensure
+          return 2
+        end
+      end
+      f3.should == 2
+    end
+
+    it 'preserves $! when an ensure runs after a return from a nested rescue' do
+      seen = nil
+      def f4
+        outer = StandardError.new 'outer'
+        begin
+          raise outer
+        rescue StandardError
+          begin
+            raise 'inner'
+          rescue StandardError
+            return $!.message
+          ensure
+            $f4_ensure_dollar_bang_message = $!.message
+          end
+        end
+      end
+      f4.should == 'inner'
+      $f4_ensure_dollar_bang_message.should == 'outer'
+    end
+  end
 end