3 files changed, 43 insertions, 7 deletions
diff --git a/lib/ruby_vm/mjit/assembler.rb b/lib/ruby_vm/mjit/assembler.rb
index a6ea501ac3..40e8827046 100644
--- a/lib/ruby_vm/mjit/assembler.rb
+++ b/lib/ruby_vm/mjit/assembler.rb
@@ -184,11 +184,31 @@ module RubyVM::MJIT
           mod_rm: ModRM[mod: Mod01, reg: right_reg, rm: left_reg],
           disp: left_disp,
         )
+      # CMP r/m64, r64 (Mod 11: reg)
+      in [Symbol => left_reg, Symbol => right_reg] if r64?(left_reg) && r64?(right_reg)
+        # REX.W + 39 /r
+        # MR: Operand 1: ModRM:r/m (r), Operand 2: ModRM:reg (r)
+        insn(
+          prefix: REX_W,
+          opcode: 0x39,
+          mod_rm: ModRM[mod: Mod11, reg: right_reg, rm: left_reg],
+        )
       else
         raise NotImplementedError, "cmp: not-implemented operands: #{left.inspect}, #{right.inspect}"
       end
     end
 
+    def jbe(dst)
+      case dst
+      # JBE rel32
+      in Integer => dst_addr
+        # 0F 86 cd
+        insn(opcode: [0x0f, 0x86], imm: rel32(dst_addr))
+      else
+        raise NotImplementedError, "jbe: not-implemented operands: #{dst.inspect}"
+      end
+    end
+
     def je(dst)
       case dst
       # JE rel32
@@ -277,7 +297,17 @@ module RubyVM::MJIT
           prefix: REX_W,
           opcode: 0x8d,
           mod_rm: ModRM[mod: Mod01, reg: dst_reg, rm: src_reg],
-          disp: src_disp,
+          disp: imm8(src_disp),
+        )
+      # LEA r64,m (Mod 10: [reg]+disp32)
+      in [Symbol => dst_reg, [Symbol => src_reg, Integer => src_disp]] if r64?(dst_reg) && r64?(src_reg) && imm32?(src_disp)
+        # REX.W + 8D /r
+        # RM: Operand 1: ModRM:reg (w), Operand 2: ModRM:r/m (r)
+        insn(
+          prefix: REX_W,
+          opcode: 0x8d,
+          mod_rm: ModRM[mod: Mod10, reg: dst_reg, rm: src_reg],
+          disp: imm32(src_disp),
         )
       else
         raise NotImplementedError, "lea: not-implemented operands: #{dst.inspect}, #{src.inspect}"
diff --git a/lib/ruby_vm/mjit/compiler.rb b/lib/ruby_vm/mjit/compiler.rb
index 4b2c5c662d..8af3c658db 100644
--- a/lib/ruby_vm/mjit/compiler.rb
+++ b/lib/ruby_vm/mjit/compiler.rb
@@ -182,10 +182,6 @@ module RubyVM::MJIT
       end
     end
 
-    def mjit_blocks(iseq)
-      iseq.body.mjit_blocks ||= {}
-    end
-
     # @param [Integer] pc
     # @param [RubyVM::MJIT::Context] ctx
     # @return [RubyVM::MJIT::Block,NilClass]
diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index 8f762ea287..da516a4476 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -1165,16 +1165,26 @@ module RubyVM::MJIT
     #
     # Frame structure:
     # | args | locals | cme/cref | block_handler/prev EP | frame type (EP here) | stack bottom (SP here)
+    #
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::Assembler]
     def jit_push_frame(jit, ctx, asm, ci, cme, flags, argc, iseq, frame_type, next_pc)
-      # TODO: stack overflow check
-
+      # CHECK_VM_STACK_OVERFLOW0: next_cfp <= sp + (local_size + stack_max) 
+      asm.comment('stack overflow check')
       local_size = iseq.body.local_table_size - iseq.body.param.size
+      asm.lea(:rax, ctx.sp_opnd(C.rb_control_frame_t.size + C.VALUE.size * (local_size + iseq.body.stack_max)))
+      asm.cmp(CFP, :rax)
+      asm.jbe(counted_exit(side_exit(jit, ctx), :send_stackoverflow))
+
       local_size.times do |i|
         asm.comment('set local variables') if i == 0
         local_index = ctx.sp_offset + i
         asm.mov([SP, C.VALUE.size * local_index], Qnil)
       end
 
+      # This moves SP register. Don't side-exit after this.
+      asm.comment('move SP register to callee stack')
       sp_offset = ctx.sp_offset + local_size + 3
       asm.add(SP, C.VALUE.size * sp_offset)