Use the actual sp_offset

4 files changed, 16 insertions, 8 deletions

diff --git a/lib/ruby_vm/mjit/assembler.rb b/lib/ruby_vm/mjit/assembler.rb
index bc42ad1a6a..f23696244d 100644
--- a/lib/ruby_vm/mjit/assembler.rb
+++ b/lib/ruby_vm/mjit/assembler.rb
@@ -215,6 +215,15 @@ module RubyVM::MJIT
             mod_rm: ModRM[mod: Mod01, reg: dst_reg, rm: src_reg],
             disp: src_disp,
           )
+        # MOV r32, imm32 (Mod 11: reg)
+        in Integer => src_imm if r32?(dst_reg) && imm32?(src_imm)
+          # B8+ rd id
+          # OI: Operand 1: opcode + rd (w), Operand 2: imm8/16/32/64
+          insn(
+            opcode: 0xb8,
+            rd: dst_reg,
+            imm: imm32(src_imm),
+          )
         # MOV r/m64, imm32 (Mod 11: reg)
         in Integer => src_imm if r64?(dst_reg) && imm32?(src_imm)
           # REX.W + C7 /0 id
diff --git a/lib/ruby_vm/mjit/exit_compiler.rb b/lib/ruby_vm/mjit/exit_compiler.rb
index d92f8a1588..a1eca9fe23 100644
--- a/lib/ruby_vm/mjit/exit_compiler.rb
+++ b/lib/ruby_vm/mjit/exit_compiler.rb
@@ -45,9 +45,10 @@ module RubyVM::MJIT
     end
 
     # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
     # @param asm [RubyVM::MJIT::Assembler]
     # @param stub [RubyVM::MJIT::BlockStub]
-    def compile_jump_stub(jit, asm, stub)
+    def compile_jump_stub(jit, ctx, asm, stub)
       case stub
       when BlockStub
         asm.comment("block stub hit: #{stub.iseq.body.location.label}@#{C.rb_iseq_path(stub.iseq)}:#{stub.iseq.body.location.first_lineno}")
@@ -57,6 +58,7 @@ module RubyVM::MJIT
 
       # Call rb_mjit_stub_hit
       asm.mov(:rdi, to_value(stub))
+      asm.mov(:esi, ctx.sp_offset)
       asm.call(C.rb_mjit_stub_hit)
 
       # Jump to the address returned by rb_mjit_stub_hit
diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index c17167a4ac..3468be9598 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -375,7 +375,7 @@ module RubyVM::MJIT
       )
 
       stub_hit = Assembler.new.then do |ocb_asm|
-        @exit_compiler.compile_jump_stub(jit, ocb_asm, block_stub)
+        @exit_compiler.compile_jump_stub(jit, ctx, ocb_asm, block_stub)
         @ocb.write(ocb_asm)
       end
 
diff --git a/mjit.c b/mjit.c
index dcf29feb84..c143bf3fa6 100644
--- a/mjit.c
+++ b/mjit.c
@@ -345,7 +345,7 @@ rb_mjit_compile(const rb_iseq_t *iseq)
 }
 
 void *
-rb_mjit_stub_hit(VALUE branch_stub)
+rb_mjit_stub_hit(VALUE branch_stub, int sp_offset)
 {
     VALUE result;
 
@@ -356,15 +356,12 @@ rb_mjit_stub_hit(VALUE branch_stub)
     mjit_stats_p = false; // Avoid impacting JIT stats by itself
 
     rb_control_frame_t *cfp = GET_EC()->cfp;
-    // Given JIT's SP offset, temporarily update SP to preserve stack values.
-    // It's reset afterwards for consistency with the code without this stub.
-    unsigned int stack_max = cfp->iseq->body->stack_max;
-    cfp->sp += stack_max;
+    cfp->sp += sp_offset; // preserve stack values, also using the actual sp_offset to make jit.peek_at_stack work
 
     VALUE cfp_ptr = rb_funcall(rb_cMJITCfpPtr, rb_intern("new"), 1, SIZET2NUM((size_t)cfp));
     result = rb_funcall(rb_MJITCompiler, rb_intern("stub_hit"), 2, branch_stub, cfp_ptr);
 
-    cfp->sp -= stack_max;
+    cfp->sp -= sp_offset; // reset for consistency with the code without the stub
 
     mjit_stats_p = mjit_opts.stats;
     mjit_call_p = original_call_p;