RJIT: Rewind stack_size on CantCompile and side exits

so that we can take an exit whenever we want.

As a starter, this commit also pops blockarg earlier than some
CantCompile exits.

4 files changed, 30 insertions, 14 deletions

diff --git a/lib/ruby_vm/rjit/compiler.rb b/lib/ruby_vm/rjit/compiler.rb
index 357fe7d734..5673926b6f 100644
--- a/lib/ruby_vm/rjit/compiler.rb
+++ b/lib/ruby_vm/rjit/compiler.rb
@@ -220,6 +220,8 @@ module RubyVM::RJIT
       while index < iseq.body.iseq_size
         insn = self.class.decode_insn(iseq.body.iseq_encoded[index])
         jit.pc = (iseq.body.iseq_encoded + index).to_i
+        jit.stack_size_for_pc = ctx.stack_size
+        jit.side_exit_for_pc.clear
 
         # If previous instruction requested to record the boundary
         if jit.record_boundary_patch_point
@@ -243,7 +245,9 @@ module RubyVM::RJIT
           # TODO: pad nops if entry exit exists (not needed for x86_64?)
           break
         when CantCompile
-          @exit_compiler.compile_side_exit(jit.pc, ctx, asm)
+          # Rewind stack_size using ctx.with_stack_size to allow stack_size changes
+          # before you return CantCompile.
+          @exit_compiler.compile_side_exit(jit.pc, ctx.with_stack_size(jit.stack_size_for_pc), asm)
 
           # If this is the first instruction, this block never needs to be invalidated.
           if block.pc == iseq.body.iseq_encoded.to_i + index * C.VALUE.size
diff --git a/lib/ruby_vm/rjit/context.rb b/lib/ruby_vm/rjit/context.rb
index 746f29b8fb..9c8fd183e6 100644
--- a/lib/ruby_vm/rjit/context.rb
+++ b/lib/ruby_vm/rjit/context.rb
@@ -26,5 +26,15 @@ module RubyVM::RJIT
     def sp_opnd(offset_bytes = 0)
       [SP, (C.VALUE.size * self.sp_offset) + offset_bytes]
     end
+
+    # Create a new Context instance with a given stack_size and sp_offset adjusted
+    # accordingly. This is useful when you want to virtually rewind a stack_size for
+    # generating a side exit while considering past sp_offset changes on gen_save_sp.
+    def with_stack_size(stack_size)
+      ctx = self.dup
+      ctx.sp_offset -= ctx.stack_size - stack_size
+      ctx.stack_size = stack_size
+      ctx
+    end
   end
 end
diff --git a/lib/ruby_vm/rjit/insn_compiler.rb b/lib/ruby_vm/rjit/insn_compiler.rb
index 62074eafaf..27031dc0e3 100644
--- a/lib/ruby_vm/rjit/insn_compiler.rb
+++ b/lib/ruby_vm/rjit/insn_compiler.rb
@@ -4237,7 +4237,10 @@ module RubyVM::RJIT
       # Number of locals that are not parameters
       num_locals = iseq.body.local_table_size - num_params
 
-      # blockarg is currently popped later
+      # Pop blockarg after all side exits
+      if flags & C::VM_CALL_ARGS_BLOCKARG != 0
+        ctx.stack_pop(1)
+      end
 
       if block_handler == C::VM_BLOCK_HANDLER_NONE && iseq.body.builtin_attrs & C::BUILTIN_ATTR_LEAF != 0
         if jit_leaf_builtin_func(jit, ctx, asm, flags, iseq)
@@ -4290,11 +4293,6 @@ module RubyVM::RJIT
         return CantCompile
       end
 
-      # Pop blockarg after all side exits
-      if flags & C::VM_CALL_ARGS_BLOCKARG != 0
-        ctx.stack_pop(1)
-      end
-
       # Setup the new frame
       frame_type ||= C::VM_FRAME_MAGIC_METHOD | C::VM_ENV_FLAG_LOCAL
       jit_push_frame(
@@ -5187,12 +5185,15 @@ module RubyVM::RJIT
     # @param jit [RubyVM::RJIT::JITState]
     # @param ctx [RubyVM::RJIT::Context]
     def side_exit(jit, ctx)
-      if side_exit = jit.side_exits[jit.pc]
-        return side_exit
+      # We use the latest ctx.sp_offset to generate a side exit to tolerate sp_offset changes by jit_save_sp.
+      # However, we want to simulate an old stack_size when we take a side exit. We do that by adjusting the
+      # sp_offset because gen_outlined_exit uses ctx.sp_offset to move SP.
+      ctx = ctx.with_stack_size(jit.stack_size_for_pc)
+
+      jit.side_exit_for_pc[ctx.sp_offset] ||= Assembler.new.then do |asm|
+        @exit_compiler.compile_side_exit(jit.pc, ctx, asm)
+        @ocb.write(asm)
       end
-      asm = Assembler.new
-      @exit_compiler.compile_side_exit(jit.pc, ctx, asm)
-      jit.side_exits[jit.pc] = @ocb.write(asm)
     end
 
     def counted_exit(side_exit, name)
diff --git a/lib/ruby_vm/rjit/jit_state.rb b/lib/ruby_vm/rjit/jit_state.rb
index a365559d17..f31223e478 100644
--- a/lib/ruby_vm/rjit/jit_state.rb
+++ b/lib/ruby_vm/rjit/jit_state.rb
@@ -4,10 +4,11 @@ module RubyVM::RJIT
     :pc,                          # @param [Integer] The JIT target PC
     :cfp,                         # @param `RubyVM::RJIT::CPointer::Struct_rb_control_frame_t` The JIT source CFP (before RJIT is called)
     :block,                       # @param [RubyVM::RJIT::Block]
-    :side_exits,                  # @param [Hash{ Integer => Integer }] { PC => address }
+    :stack_size_for_pc,           # @param [Integer]
+    :side_exit_for_pc,            # @param [Hash{ Integer => Integer }] { sp_offset => address }
     :record_boundary_patch_point, # @param [TrueClass,FalseClass]
   )
-    def initialize(side_exits: {}, record_boundary_patch_point: false, **) = super
+    def initialize(side_exit_for_pc: {}, record_boundary_patch_point: false, **) = super
 
     def insn
       Compiler.decode_insn(C.VALUE.new(pc).*)