Implement invalidation after cfunc

8 files changed, 143 insertions, 78 deletions

diff --git a/lib/ruby_vm/mjit/assembler.rb b/lib/ruby_vm/mjit/assembler.rb
index f8543c4a66..e50cbdf600 100644
--- a/lib/ruby_vm/mjit/assembler.rb
+++ b/lib/ruby_vm/mjit/assembler.rb
@@ -36,6 +36,7 @@ module RubyVM::MJIT
       @blocks = Hash.new { |h, k| h[k] = [] }
       @stub_starts = Hash.new { |h, k| h[k] = [] }
       @stub_ends = Hash.new { |h, k| h[k] = [] }
+      @pos_markers = Hash.new { |h, k| h[k] = [] }
     end
 
     def assemble(addr)
@@ -45,6 +46,9 @@ module RubyVM::MJIT
 
       write_bytes(addr)
 
+      @pos_markers.each do |write_pos, markers|
+        markers.each { |marker| marker.call(addr + write_pos) }
+      end
       @bytes.size
     ensure
       @bytes.clear
@@ -617,6 +621,10 @@ module RubyVM::MJIT
       @stub_ends[@bytes.size] << stub
     end
 
+    def pos_marker(&block)
+      @pos_markers[@bytes.size] << block
+    end
+
     def new_label(name)
       Label.new(id: @label_id += 1, name:)
     end
diff --git a/lib/ruby_vm/mjit/compiler.rb b/lib/ruby_vm/mjit/compiler.rb
index 8af3c658db..4f5779b593 100644
--- a/lib/ruby_vm/mjit/compiler.rb
+++ b/lib/ruby_vm/mjit/compiler.rb
@@ -45,6 +45,7 @@ module RubyVM::MJIT
       @ocb = CodeBlock.new(mem_block: mem_block + mem_size / 2, mem_size: mem_size / 2, outlined: true)
       @exit_compiler = ExitCompiler.new
       @insn_compiler = InsnCompiler.new(@cb, @ocb, @exit_compiler)
+      Invariants.initialize(@cb, @ocb, @exit_compiler)
 
       @leave_exit = Assembler.new.then do |asm|
         @exit_compiler.compile_leave_exit(asm)
diff --git a/lib/ruby_vm/mjit/exit_compiler.rb b/lib/ruby_vm/mjit/exit_compiler.rb
index ef12740153..531fe3f426 100644
--- a/lib/ruby_vm/mjit/exit_compiler.rb
+++ b/lib/ruby_vm/mjit/exit_compiler.rb
@@ -24,7 +24,8 @@ module RubyVM::MJIT
       asm.ret
     end
 
-    # @param ocb [CodeBlock]
+    # Set to cfp->jit_return by default for leave insn
+    # @param asm [RubyVM::MJIT::Assembler]
     def compile_leave_exit(asm)
       asm.comment('default cfp->jit_return')
 
@@ -37,6 +38,28 @@ module RubyVM::MJIT
       asm.ret
     end
 
+    # Fire cfunc events on invalidation by TracePoint
+    # @param asm [RubyVM::MJIT::Assembler]
+    def compile_full_cfunc_return(asm)
+      # This chunk of code expects REG_EC to be filled properly and
+      # RAX to contain the return value of the C method.
+
+      asm.comment('full cfunc return')
+      asm.mov(C_ARG_OPNDS[0], EC)
+      asm.mov(C_ARG_OPNDS[1], :rax)
+      asm.call(C.rb_full_cfunc_return)
+
+      # TODO: count the exit
+
+      # Restore callee-saved registers
+      asm.pop(SP)
+      asm.pop(EC)
+      asm.pop(CFP)
+
+      asm.mov(:rax, Qundef)
+      asm.ret
+    end
+
     # @param jit [RubyVM::MJIT::JITState]
     # @param ctx [RubyVM::MJIT::Context]
     # @param asm [RubyVM::MJIT::Assembler]
diff --git a/lib/ruby_vm/mjit/hooks.rb b/lib/ruby_vm/mjit/hooks.rb
index 6aef14b56a..39a949b4d7 100644
--- a/lib/ruby_vm/mjit/hooks.rb
+++ b/lib/ruby_vm/mjit/hooks.rb
@@ -1,32 +1,27 @@
-module RubyVM::MJIT::Hooks # :nodoc: all
-  C = RubyVM::MJIT.const_get(:C, false)
+module RubyVM::MJIT
+  module Hooks # :nodoc: all
+    def self.on_bop_redefined(_redefined_flag, _bop)
+      # C.mjit_cancel_all("BOP is redefined")
+    end
 
-  def self.on_bop_redefined(_redefined_flag, _bop)
-    # C.mjit_cancel_all("BOP is redefined")
-  end
-
-  def self.on_cme_invalidate(_cme)
-    # to be used later
-  end
+    def self.on_cme_invalidate(cme)
+      Invariants.on_cme_invalidate(cme)
+    end
 
-  def self.on_ractor_spawn
-    # C.mjit_cancel_all("Ractor is spawned")
-  end
+    def self.on_ractor_spawn
+      # C.mjit_cancel_all("Ractor is spawned")
+    end
 
-  def self.on_constant_state_changed(_id)
-    # to be used later
-  end
+    def self.on_constant_state_changed(_id)
+      # to be used later
+    end
 
-  def self.on_constant_ic_update(_iseq, _ic, _insn_idx)
-    # to be used later
-  end
+    def self.on_constant_ic_update(_iseq, _ic, _insn_idx)
+      # to be used later
+    end
 
-  def self.on_tracing_invalidate_all(new_iseq_events)
-    # # Stop calling all JIT-ed code. We can't rewrite existing JIT-ed code to trace_ insns for now.
-    # # :class events are triggered only in ISEQ_TYPE_CLASS, but mjit_target_iseq_p ignores such iseqs.
-    # # Thus we don't need to cancel JIT-ed code for :class events.
-    # if new_iseq_events != C.RUBY_EVENT_CLASS
-    #   C.mjit_cancel_all("TracePoint is enabled")
-    # end
+    def self.on_tracing_invalidate_all(_new_iseq_events)
+      Invariants.on_tracing_invalidate_all
+    end
   end
 end
diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index 139d0851d0..6898c779e1 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -5,8 +5,13 @@ module RubyVM::MJIT
     def initialize(cb, ocb, exit_compiler)
       @ocb = ocb
       @exit_compiler = exit_compiler
-      @invariants = Invariants.new(cb, ocb, exit_compiler)
       @gc_refs = [] # TODO: GC offsets?
+
+      @full_cfunc_return = Assembler.new.then do |asm|
+        @exit_compiler.compile_full_cfunc_return(asm)
+        @ocb.write(asm)
+      end
+
       # freeze # workaround a binding.irb issue. TODO: resurrect this
     end
 
@@ -421,7 +426,7 @@ module RubyVM::MJIT
         # Generate a side exit before popping operands
         side_exit = side_exit(jit, ctx)
 
-        unless @invariants.assume_bop_not_redefined(jit, C.INTEGER_REDEFINED_OP_FLAG, C.BOP_PLUS)
+        unless Invariants.assume_bop_not_redefined(jit, C.INTEGER_REDEFINED_OP_FLAG, C.BOP_PLUS)
           return CantCompile
         end
 
@@ -467,7 +472,7 @@ module RubyVM::MJIT
         # Generate a side exit before popping operands
         side_exit = side_exit(jit, ctx)
 
-        unless @invariants.assume_bop_not_redefined(jit, C.INTEGER_REDEFINED_OP_FLAG, C.BOP_MINUS)
+        unless Invariants.assume_bop_not_redefined(jit, C.INTEGER_REDEFINED_OP_FLAG, C.BOP_MINUS)
           return CantCompile
         end
 
@@ -531,7 +536,7 @@ module RubyVM::MJIT
         # Generate a side exit before popping operands
         side_exit = side_exit(jit, ctx)
 
-        unless @invariants.assume_bop_not_redefined(jit, C.INTEGER_REDEFINED_OP_FLAG, C.BOP_LT)
+        unless Invariants.assume_bop_not_redefined(jit, C.INTEGER_REDEFINED_OP_FLAG, C.BOP_LT)
           return CantCompile
         end
 
@@ -601,7 +606,7 @@ module RubyVM::MJIT
         asm.incr_counter(:optaref_array)
         CantCompile
       elsif comptime_recv.class == Hash
-        unless @invariants.assume_bop_not_redefined(jit, C.HASH_REDEFINED_OP_FLAG, C.BOP_AREF)
+        unless Invariants.assume_bop_not_redefined(jit, C.HASH_REDEFINED_OP_FLAG, C.BOP_AREF)
           return CantCompile
         end
 
@@ -1051,7 +1056,7 @@ module RubyVM::MJIT
       end
 
       # Invalidate on redefinition (part of vm_search_method_fastpath)
-      @invariants.assume_method_lookup_stable(jit, cme)
+      Invariants.assume_method_lookup_stable(jit, cme)
 
       jit_call_method_each_type(jit, ctx, asm, ci, argc, flags, cme, comptime_recv, recv_opnd)
     end
@@ -1155,7 +1160,7 @@ module RubyVM::MJIT
         return CantCompile
       end
 
-      # Disabled until we implement TracePoint invalidation
+      # Disabled until we figure out why $' gets broken on test-all
       disabled = true
       if disabled
         return CantCompile
@@ -1214,6 +1219,8 @@ module RubyVM::MJIT
       asm.call(:rax) # TODO: use rel32 if close enough
       ctx.stack_pop(1 + argc)
 
+      Invariants.record_global_inval_patch(asm, @full_cfunc_return)
+
       asm.comment('push the return value')
       stack_ret = ctx.stack_push
       asm.mov(stack_ret, :rax)
diff --git a/lib/ruby_vm/mjit/invariants.rb b/lib/ruby_vm/mjit/invariants.rb
index 314fe90b8b..6dd8e65bc9 100644
--- a/lib/ruby_vm/mjit/invariants.rb
+++ b/lib/ruby_vm/mjit/invariants.rb
@@ -2,61 +2,82 @@ require 'set'
 
 module RubyVM::MJIT
   class Invariants
-    # @param cb [CodeBlock]
-    # @param ocb [CodeBlock]
-    # @param exit_compiler [RubyVM::MJIT::ExitCompiler]
-    def initialize(cb, ocb, exit_compiler)
-      @cb = cb
-      @ocb = ocb
-      @exit_compiler = exit_compiler
-      @bop_blocks = Set.new # TODO: actually invalidate this
-      @cme_blocks = Hash.new { |h, k| h[k] = Set.new }
+    class << self
+      # Called by RubyVM::MJIT::Compiler to lazily initialize this
+      # @param cb [CodeBlock]
+      # @param ocb [CodeBlock]
+      # @param exit_compiler [RubyVM::MJIT::ExitCompiler]
+      def initialize(cb, ocb, exit_compiler)
+        @cb = cb
+        @ocb = ocb
+        @exit_compiler = exit_compiler
+        @bop_blocks = Set.new # TODO: actually invalidate this
+        @cme_blocks = Hash.new { |h, k| h[k] = Set.new }
+        @patches = {}
 
-      invariants = self
-      hooks = Module.new
-      hooks.define_method(:on_cme_invalidate) do |cme|
-        invariants.on_cme_invalidate(cme)
+        # freeze # workaround a binding.irb issue. TODO: resurrect this
       end
-      Hooks.singleton_class.prepend(hooks)
-    end
 
-    # @param jit [RubyVM::MJIT::JITState]
-    # @param klass [Integer]
-    # @param op [Integer]
-    def assume_bop_not_redefined(jit, klass, op)
-      return false unless C.BASIC_OP_UNREDEFINED_P(klass, op)
+      # @param jit [RubyVM::MJIT::JITState]
+      # @param klass [Integer]
+      # @param op [Integer]
+      def assume_bop_not_redefined(jit, klass, op)
+        return false unless C.BASIC_OP_UNREDEFINED_P(klass, op)
 
-      ensure_block_entry_exit(jit.block, cause: 'assume_bop_not_redefined')
-      @bop_blocks << jit.block
-      true
-    end
+        ensure_block_entry_exit(jit.block, cause: 'assume_bop_not_redefined')
+        @bop_blocks << jit.block
+        true
+      end
 
-    # @param jit [RubyVM::MJIT::JITState]
-    def assume_method_lookup_stable(jit, cme)
-      ensure_block_entry_exit(jit.block, cause: 'assume_method_lookup_stable')
-      @cme_blocks[cme.to_i] << jit.block
-    end
+      # @param jit [RubyVM::MJIT::JITState]
+      def assume_method_lookup_stable(jit, cme)
+        ensure_block_entry_exit(jit.block, cause: 'assume_method_lookup_stable')
+        @cme_blocks[cme.to_i] << jit.block
+      end
 
-    def on_cme_invalidate(cme)
-      @cme_blocks.fetch(cme.to_i, []).each do |block|
-        @cb.with_write_addr(block.start_addr) do
-          asm = Assembler.new
-          asm.comment('on_cme_invalidate')
-          asm.jmp(block.entry_exit)
-          @cb.write(asm)
+      # @param asm [RubyVM::MJIT::Assembler]
+      def record_global_inval_patch(asm, target)
+        asm.pos_marker do |address|
+          if @patches.key?(address)
+            raise 'multiple patches in the same address'
+          end
+          @patches[address] = target
+        end
+      end
+
+      def on_cme_invalidate(cme)
+        @cme_blocks.fetch(cme.to_i, []).each do |block|
+          @cb.with_write_addr(block.start_addr) do
+            asm = Assembler.new
+            asm.comment('on_cme_invalidate')
+            asm.jmp(block.entry_exit)
+            @cb.write(asm)
+          end
+          # TODO: re-generate branches that refer to this block
+        end
+      end
+
+      def on_tracing_invalidate_all
+        # TODO: assert patches don't overlap each other
+        @patches.each do |address, target|
+          @cb.with_write_addr(address) do
+            asm = Assembler.new
+            asm.comment('on_tracing_invalidate_all')
+            asm.jmp(target)
+            @cb.write(asm)
+          end
         end
-        # TODO: re-generate branches that refer to this block
       end
-    end
 
-    private
+      private
 
-    # @param block [RubyVM::MJIT::Block]
-    def ensure_block_entry_exit(block, cause:)
-      if block.entry_exit.nil?
-        block.entry_exit = Assembler.new.then do |asm|
-          @exit_compiler.compile_entry_exit(block.pc, block.ctx, asm, cause:)
-          @ocb.write(asm)
+      # @param block [RubyVM::MJIT::Block]
+      def ensure_block_entry_exit(block, cause:)
+        if block.entry_exit.nil?
+          block.entry_exit = Assembler.new.then do |asm|
+            @exit_compiler.compile_entry_exit(block.pc, block.ctx, asm, cause:)
+            @ocb.write(asm)
+          end
         end
       end
     end
diff --git a/mjit.c b/mjit.c
index 607ec5c59c..fc19896ff6 100644
--- a/mjit.c
+++ b/mjit.c
@@ -353,6 +353,9 @@ void
 rb_mjit_tracing_invalidate_all(rb_event_flag_t new_iseq_events)
 {
     if (!mjit_call_p) return;
+    WITH_MJIT_DISABLED({
+        rb_funcall(rb_mMJITHooks, rb_intern("on_tracing_invalidate_all"), 1, UINT2NUM(new_iseq_events));
+    });
     mjit_call_p = false;
 }
 
diff --git a/mjit_c.rb b/mjit_c.rb
index 479c6d0abf..bee063aa3e 100644
--- a/mjit_c.rb
+++ b/mjit_c.rb
@@ -138,6 +138,13 @@ module RubyVM::MJIT # :nodoc: all
       }
     end
 
+    def rb_full_cfunc_return
+      Primitive.cstmt! %{
+        extern void rb_full_cfunc_return(rb_execution_context_t *ec, VALUE return_value);
+        return SIZET2NUM((size_t)rb_full_cfunc_return);
+      }
+    end
+
     #========================================================================================
     #
     # Old stuff