RJIT: Support entry with different PCs

4 files changed, 74 insertions, 19 deletions

diff --git a/lib/ruby_vm/rjit/compiler.rb b/lib/ruby_vm/rjit/compiler.rb
index 98ce556357..39b3c77f3e 100644
--- a/lib/ruby_vm/rjit/compiler.rb
+++ b/lib/ruby_vm/rjit/compiler.rb
@@ -3,6 +3,7 @@ require 'ruby_vm/rjit/block'
 require 'ruby_vm/rjit/branch_stub'
 require 'ruby_vm/rjit/code_block'
 require 'ruby_vm/rjit/context'
+require 'ruby_vm/rjit/entry_stub'
 require 'ruby_vm/rjit/exit_compiler'
 require 'ruby_vm/rjit/insn_compiler'
 require 'ruby_vm/rjit/instruction'
@@ -64,6 +65,48 @@ module RubyVM::RJIT
       exit 1
     end
 
+    # Compile an entry.
+    # @param entry [RubyVM::RJIT::EntryStub]
+    def entry_stub_hit(entry_stub, cfp)
+      # Compile a new entry guard as a next entry
+      pc = cfp.pc.to_i
+      next_entry = Assembler.new.then do |asm|
+        compile_entry_chain_guard(asm, cfp.iseq, pc)
+        @cb.write(asm)
+      end
+
+      # Try to find an existing compiled version of this block
+      ctx = Context.new
+      block = find_block(cfp.iseq, pc, ctx)
+      if block
+        # If an existing block is found, generate a jump to the block.
+        asm = Assembler.new
+        asm.jmp(block.start_addr)
+        @cb.write(asm)
+      else
+        # If this block hasn't yet been compiled, generate blocks after the entry guard.
+        asm = Assembler.new
+        jit = JITState.new(iseq: cfp.iseq, cfp:)
+        compile_block(asm, jit:, pc:, ctx:)
+        @cb.write(asm)
+
+        block = jit.block
+      end
+
+      # Regenerate the previous entry
+      @cb.with_write_addr(entry_stub.start_addr) do
+        # The last instruction of compile_entry_chain_guard is jne
+        asm = Assembler.new
+        asm.jne(next_entry)
+        @cb.write(asm)
+      end
+
+      return block.start_addr
+    rescue Exception => e
+      $stderr.puts e.full_message
+      exit 1
+    end
+
     # Compile a branch stub.
     # @param branch_stub [RubyVM::RJIT::BranchStub]
     # @param cfp `RubyVM::RJIT::CPointer::Struct_rb_control_frame_t`
@@ -210,30 +253,24 @@ module RubyVM::RJIT
       # 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)
+        compile_entry_chain_guard(asm, iseq, pc)
       end
     end
 
-    def compile_pc_guard(asm, iseq, pc)
+    def compile_entry_chain_guard(asm, iseq, pc)
+      entry_stub = EntryStub.new
+      stub_addr = Assembler.new.then do |ocb_asm|
+        @exit_compiler.compile_entry_stub(ocb_asm, entry_stub)
+        @ocb.write(ocb_asm)
+      end
+
       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)
+      asm.stub(entry_stub) do
+        asm.jne(stub_addr)
+      end
     end
 
     # @param asm [RubyVM::RJIT::Assembler]
diff --git a/lib/ruby_vm/rjit/entry_stub.rb b/lib/ruby_vm/rjit/entry_stub.rb
new file mode 100644
index 0000000000..9bcef14053
--- /dev/null
+++ b/lib/ruby_vm/rjit/entry_stub.rb
@@ -0,0 +1,7 @@
+module RubyVM::RJIT
+  class EntryStub < Struct.new(
+    :start_addr, # @param [Integer] Stub source start address to be re-generated
+    :end_addr,   # @param [Integer] Stub source end address to be re-generated
+  )
+  end
+end
diff --git a/lib/ruby_vm/rjit/exit_compiler.rb b/lib/ruby_vm/rjit/exit_compiler.rb
index b7beb22177..1ced2141a4 100644
--- a/lib/ruby_vm/rjit/exit_compiler.rb
+++ b/lib/ruby_vm/rjit/exit_compiler.rb
@@ -78,6 +78,18 @@ module RubyVM::RJIT
       asm.ret
     end
 
+    # @param asm [RubyVM::RJIT::Assembler]
+    # @param entry_stub [RubyVM::RJIT::EntryStub]
+    def compile_entry_stub(asm, entry_stub)
+      # Call rb_rjit_entry_stub_hit
+      asm.comment('entry stub hit')
+      asm.mov(C_ARGS[0], to_value(entry_stub))
+      asm.call(C.rb_rjit_entry_stub_hit)
+
+      # Jump to the address returned by rb_rjit_entry_stub_hit
+      asm.jmp(:rax)
+    end
+
     # @param ctx [RubyVM::RJIT::Context]
     # @param asm [RubyVM::RJIT::Assembler]
     # @param branch_stub [RubyVM::RJIT::BranchStub]
@@ -93,7 +105,7 @@ module RubyVM::RJIT
       asm.mov(:edx, target0_p ? 1 : 0)
       asm.call(C.rb_rjit_branch_stub_hit)
 
-      # Jump to the address returned by rb_rjit_stub_hit
+      # Jump to the address returned by rb_rjit_branch_stub_hit
       asm.jmp(:rax)
     end
 
diff --git a/lib/ruby_vm/rjit/stats.rb b/lib/ruby_vm/rjit/stats.rb
index 778b68d19a..8c4253880a 100644
--- a/lib/ruby_vm/rjit/stats.rb
+++ b/lib/ruby_vm/rjit/stats.rb
@@ -40,7 +40,6 @@ 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')