RJIT: Support has_opt ISEQs

4 files changed, 43 insertions, 7 deletions

diff --git a/lib/ruby_vm/rjit/compiler.rb b/lib/ruby_vm/rjit/compiler.rb
index a7038b6d2b..98ce556357 100644
--- a/lib/ruby_vm/rjit/compiler.rb
+++ b/lib/ruby_vm/rjit/compiler.rb
@@ -53,13 +53,11 @@ module RubyVM::RJIT
     # @param iseq `RubyVM::RJIT::CPointer::Struct_rb_iseq_t`
     # @param cfp `RubyVM::RJIT::CPointer::Struct_rb_control_frame_t`
     def compile(iseq, cfp)
-      # TODO: Support has_opt
-      return if iseq.body.param.flags.has_opt
-
+      pc = cfp.pc.to_i
       jit = JITState.new(iseq:, cfp:)
       asm = Assembler.new
-      compile_prologue(asm)
-      compile_block(asm, jit:)
+      compile_prologue(asm, iseq, pc)
+      compile_block(asm, jit:, pc:)
       iseq.body.jit_func = @cb.write(asm)
     rescue Exception => e
       $stderr.puts e.full_message
@@ -186,7 +184,7 @@ module RubyVM::RJIT
     # Caller-saved: rax, rdi, rsi, rdx, rcx, r8, r9, r10, r11
     #
     # @param asm [RubyVM::RJIT::Assembler]
-    def compile_prologue(asm)
+    def compile_prologue(asm, iseq, pc)
       asm.comment('RJIT entry point')
 
       # Save callee-saved registers used by JITed code
@@ -204,10 +202,42 @@ module RubyVM::RJIT
       # Setup cfp->jit_return
       asm.mov(:rax, leave_exit)
       asm.mov([CFP, C.rb_control_frame_t.offsetof(:jit_return)], :rax)
+
+      # We're compiling iseqs that we *expect* to start at `insn_idx`. But in
+      # the case of optional parameters, the interpreter can set the pc to a
+      # different location depending on the optional parameters.  If an iseq
+      # has optional parameters, we'll add a runtime check that the PC we've
+      # compiled for is the same PC that the interpreter wants us to run with.
+      # If they don't match, then we'll take a side exit.
+      if iseq.body.param.flags.has_opt
+        compile_pc_guard(asm, iseq, pc)
+      end
+    end
+
+    def compile_pc_guard(asm, iseq, pc)
+      asm.comment('guard expected PC')
+      asm.mov(:rax, pc)
+      asm.cmp([CFP, C.rb_control_frame_t.offsetof(:pc)], :rax)
+
+      pc_match = asm.new_label('pc_match')
+      asm.je(pc_match)
+
+      # We're not starting at the first PC, so we need to exit.
+      asm.incr_counter(:leave_start_pc_non_zero)
+
+      asm.pop(SP)
+      asm.pop(EC)
+      asm.pop(CFP)
+
+      asm.mov(:rax, Qundef)
+      asm.ret
+
+      # PC should match the expected insn_idx
+      asm.write_label(pc_match)
     end
 
     # @param asm [RubyVM::RJIT::Assembler]
-    def compile_block(asm, jit:, pc: jit.iseq.body.iseq_encoded.to_i, ctx: Context.new)
+    def compile_block(asm, jit:, pc:, ctx: Context.new)
       # Mark the block start address and prepare an exit code storage
       block = Block.new(iseq: jit.iseq, pc:, ctx: ctx.dup)
       jit.block = block
diff --git a/lib/ruby_vm/rjit/stats.rb b/lib/ruby_vm/rjit/stats.rb
index 8c4253880a..778b68d19a 100644
--- a/lib/ruby_vm/rjit/stats.rb
+++ b/lib/ruby_vm/rjit/stats.rb
@@ -40,6 +40,7 @@ module RubyVM::RJIT
       print_counters(stats, prefix: 'send_', prompt: 'method call exit reasons')
       print_counters(stats, prefix: 'invokeblock_', prompt: 'invokeblock exit reasons')
       print_counters(stats, prefix: 'invokesuper_', prompt: 'invokesuper exit reasons')
+      print_counters(stats, prefix: 'leave_', prompt: 'leave exit reasons')
       print_counters(stats, prefix: 'getblockpp_', prompt: 'getblockparamproxy exit reasons')
       print_counters(stats, prefix: 'getivar_', prompt: 'getinstancevariable exit reasons')
       print_counters(stats, prefix: 'setivar_', prompt: 'setinstancevariable exit reasons')
diff --git a/rjit_c.h b/rjit_c.h
index 5756ec19ef..84b86734ef 100644
--- a/rjit_c.h
+++ b/rjit_c.h
@@ -78,6 +78,7 @@ RJIT_RUNTIME_COUNTERS(
     send_iseq_kwargs_mismatch,
     send_iseq_splat_with_kw,
     send_iseq_splat_arity_error,
+    send_iseq_has_rest_and_splat_not_equal,
 
     send_cfunc_variadic,
     send_cfunc_too_many_args,
@@ -156,6 +157,8 @@ RJIT_RUNTIME_COUNTERS(
     getblockpp_not_gc_guarded,
     getblockpp_not_iseq_block,
 
+    leave_start_pc_non_zero,
+
     compiled_block_count
 )
 #undef RJIT_RUNTIME_COUNTERS
diff --git a/rjit_c.rb b/rjit_c.rb
index bd7ef3bf0b..f873b4c89b 100644
--- a/rjit_c.rb
+++ b/rjit_c.rb
@@ -1372,6 +1372,7 @@ module RubyVM::RJIT # :nodoc: all
       send_iseq_kwargs_mismatch: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_iseq_kwargs_mismatch)")],
       send_iseq_splat_with_kw: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_iseq_splat_with_kw)")],
       send_iseq_splat_arity_error: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_iseq_splat_arity_error)")],
+      send_iseq_has_rest_and_splat_not_equal: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_iseq_has_rest_and_splat_not_equal)")],
       send_cfunc_variadic: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_cfunc_variadic)")],
       send_cfunc_too_many_args: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_cfunc_too_many_args)")],
       send_cfunc_ruby_array_varg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_cfunc_ruby_array_varg)")],
@@ -1435,6 +1436,7 @@ module RubyVM::RJIT # :nodoc: all
       getblockpp_block_handler_none: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), getblockpp_block_handler_none)")],
       getblockpp_not_gc_guarded: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), getblockpp_not_gc_guarded)")],
       getblockpp_not_iseq_block: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), getblockpp_not_iseq_block)")],
+      leave_start_pc_non_zero: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), leave_start_pc_non_zero)")],
       compiled_block_count: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), compiled_block_count)")],
     )
   end