Invalidate blocks on global constant changes

7 files changed, 88 insertions, 46 deletions

diff --git a/lib/ruby_vm/mjit/block.rb b/lib/ruby_vm/mjit/block.rb
index abe0c221b2..9b0c91ad40 100644
--- a/lib/ruby_vm/mjit/block.rb
+++ b/lib/ruby_vm/mjit/block.rb
@@ -1,4 +1,5 @@
 class RubyVM::MJIT::Block < Struct.new(
+  :iseq,        # @param ``
   :pc,          # @param [Integer] Starting PC
   :ctx,         # @param [RubyVM::MJIT::Context] **Starting** Context (TODO: freeze?)
   :start_addr,  # @param [Integer] Starting address of this block's JIT code
diff --git a/lib/ruby_vm/mjit/compiler.rb b/lib/ruby_vm/mjit/compiler.rb
index 30f2557f5b..892b250895 100644
--- a/lib/ruby_vm/mjit/compiler.rb
+++ b/lib/ruby_vm/mjit/compiler.rb
@@ -131,7 +131,7 @@ module RubyVM::MJIT
     # @param pc [Integer]
     def invalidate_blocks(iseq, pc)
       list_blocks(iseq, pc).each do |block|
-        invalidate_block(iseq, block)
+        invalidate_block(block)
       end
 
       # If they were the ISEQ's first blocks, re-compile MJIT entry as well
@@ -141,6 +141,48 @@ module RubyVM::MJIT
       end
     end
 
+    def invalidate_block(block)
+      iseq = block.iseq
+      # Remove this block from the version array
+      remove_block(iseq, block)
+
+      # Invalidate the block with entry exit
+      unless block.invalidated
+        @cb.with_write_addr(block.start_addr) do
+          asm = Assembler.new
+          asm.comment('invalidate_block')
+          asm.jmp(block.entry_exit)
+          @cb.write(asm)
+        end
+        block.invalidated = true
+      end
+
+      # Re-stub incoming branches
+      block.incoming.each do |branch_stub|
+        target = [branch_stub.target0, branch_stub.target1].compact.find do |target|
+          target.pc == block.pc && target.ctx == block.ctx
+        end
+        next if target.nil?
+        # TODO: Could target.address be a stub address? Is invalidation not needed in that case?
+
+        # If the target being re-generated is currently a fallthrough block,
+        # the fallthrough code must be rewritten with a jump to the stub.
+        if target.address == branch_stub.end_addr
+          branch_stub.shape = Default
+        end
+
+        target.address = Assembler.new.then do |ocb_asm|
+          @exit_compiler.compile_branch_stub(block.ctx, ocb_asm, branch_stub, target == branch_stub.target0)
+          @ocb.write(ocb_asm)
+        end
+        @cb.with_write_addr(branch_stub.start_addr) do
+          branch_asm = Assembler.new
+          branch_stub.compile.call(branch_asm)
+          @cb.write(branch_asm)
+        end
+      end
+    end
+
     private
 
     # Callee-saved: rbx, rsp, rbp, r12, r13, r14, r15
@@ -170,7 +212,7 @@ module RubyVM::MJIT
     # @param asm [RubyVM::MJIT::Assembler]
     def compile_block(asm, jit:, pc: jit.iseq.body.iseq_encoded.to_i, ctx: Context.new)
       # Mark the block start address and prepare an exit code storage
-      block = Block.new(pc:, ctx: ctx.dup)
+      block = Block.new(iseq: jit.iseq, pc:, ctx: ctx.dup)
       jit.block = block
       asm.block(block)
 
@@ -222,47 +264,6 @@ module RubyVM::MJIT
       end
     end
 
-    def invalidate_block(iseq, block)
-      # Remove this block from the version array
-      remove_block(iseq, block)
-
-      # Invalidate the block with entry exit
-      unless block.invalidated
-        @cb.with_write_addr(block.start_addr) do
-          asm = Assembler.new
-          asm.comment('invalidate_block')
-          asm.jmp(block.entry_exit)
-          @cb.write(asm)
-        end
-        block.invalidated = true
-      end
-
-      # Re-stub incoming branches
-      block.incoming.each do |branch_stub|
-        target = [branch_stub.target0, branch_stub.target1].compact.find do |target|
-          target.pc == block.pc && target.ctx == block.ctx
-        end
-        next if target.nil?
-        # TODO: Could target.address be a stub address? Is invalidation not needed in that case?
-
-        # If the target being re-generated is currently a fallthrough block,
-        # the fallthrough code must be rewritten with a jump to the stub.
-        if target.address == branch_stub.end_addr
-          branch_stub.shape = Default
-        end
-
-        target.address = Assembler.new.then do |ocb_asm|
-          @exit_compiler.compile_branch_stub(block.ctx, ocb_asm, branch_stub, target == branch_stub.target0)
-          @ocb.write(ocb_asm)
-        end
-        @cb.with_write_addr(branch_stub.start_addr) do
-          branch_asm = Assembler.new
-          branch_stub.compile.call(branch_asm)
-          @cb.write(branch_asm)
-        end
-      end
-    end
-
     def list_blocks(iseq, pc)
       mjit_blocks(iseq)[pc].values
     end
diff --git a/lib/ruby_vm/mjit/hooks.rb b/lib/ruby_vm/mjit/hooks.rb
index 477c4e0ae8..57d4ebc381 100644
--- a/lib/ruby_vm/mjit/hooks.rb
+++ b/lib/ruby_vm/mjit/hooks.rb
@@ -13,10 +13,12 @@ module RubyVM::MJIT
       # C.mjit_cancel_all("Ractor is spawned")
     end
 
-    def self.on_constant_state_changed(_id)
-      # to be used later
+    # Global constant changes like const_set
+    def self.on_constant_state_changed(id)
+      Invariants.on_constant_state_changed(id)
     end
 
+    # ISEQ-specific constant invalidation
     def self.on_constant_ic_update(iseq, ic, insn_idx)
       iseq = C.rb_iseq_t.new(iseq)
       ic = C.IC.new(ic)
diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index e3dff5e86c..756e04df4c 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -249,7 +249,7 @@ module RubyVM::MJIT
 
         # Invalidate output code on any constant writes associated with
         # constants referenced within the current block.
-        #assume_stable_constant_names(jit, ocb, idlist);
+        Invariants.assume_stable_constant_names(jit, idlist)
 
         putobject(jit, ctx, asm, val: ice.value)
       end
diff --git a/lib/ruby_vm/mjit/invariants.rb b/lib/ruby_vm/mjit/invariants.rb
index 648044f027..127e446ed9 100644
--- a/lib/ruby_vm/mjit/invariants.rb
+++ b/lib/ruby_vm/mjit/invariants.rb
@@ -15,6 +15,7 @@ module RubyVM::MJIT
         @exit_compiler = exit_compiler
         @bop_blocks = Set.new # TODO: actually invalidate this
         @cme_blocks = Hash.new { |h, k| h[k] = Set.new }
+        @const_blocks = Hash.new { |h, k| h[k] = Set.new }
         @patches = {}
 
         # freeze # workaround a binding.irb issue. TODO: resurrect this
@@ -37,6 +38,13 @@ module RubyVM::MJIT
         @cme_blocks[cme.to_i] << jit.block
       end
 
+      def assume_stable_constant_names(jit, idlist)
+        (0..).each do |i|
+          break if (id = idlist[i]) == 0
+          @const_blocks[id] << jit.block
+        end
+      end
+
       # @param asm [RubyVM::MJIT::Assembler]
       def record_global_inval_patch(asm, target)
         asm.pos_marker do |address|
@@ -57,6 +65,7 @@ module RubyVM::MJIT
           end
           # TODO: re-generate branches that refer to this block
         end
+        @cme_blocks.delete(cme.to_i)
       end
 
       def on_constant_ic_update(iseq, ic, insn_idx)
@@ -77,6 +86,12 @@ module RubyVM::MJIT
         @compiler.invalidate_blocks(iseq, pc.to_i)
       end
 
+      def on_constant_state_changed(id)
+        @const_blocks.fetch(id, []).each do |block|
+          @compiler.invalidate_block(block)
+        end
+      end
+
       def on_tracing_invalidate_all
         invalidate_all
       end
diff --git a/mjit.c b/mjit.c
index b9a5b025ce..fe5e2976a7 100644
--- a/mjit.c
+++ b/mjit.c
@@ -353,6 +353,28 @@ rb_mjit_before_ractor_spawn(void)
     mjit_call_p = false;
 }
 
+static void
+mjit_constant_state_changed(void *data)
+{
+    if (!mjit_enabled || !mjit_call_p || !rb_mMJITHooks) return;
+    RB_VM_LOCK_ENTER();
+    rb_vm_barrier();
+
+    WITH_MJIT_ISOLATED({
+        rb_funcall(rb_mMJITHooks, rb_intern("on_constant_state_changed"), 1, SIZET2NUM((size_t)data));
+    });
+
+    RB_VM_LOCK_LEAVE();
+}
+
+void
+rb_mjit_constant_state_changed(ID id)
+{
+    if (!mjit_enabled || !mjit_call_p || !rb_mMJITHooks) return;
+    // Asynchronously hook the Ruby code since this is hooked during a "Ruby critical section".
+    rb_workqueue_register(0, mjit_constant_state_changed, (void *)id);
+}
+
 void
 rb_mjit_constant_ic_update(const rb_iseq_t *const iseq, IC ic, unsigned insn_idx)
 {
diff --git a/vm_method.c b/vm_method.c
index 4cb96cc7fd..12947e1aa5 100644
--- a/vm_method.c
+++ b/vm_method.c
@@ -150,6 +150,7 @@ rb_clear_constant_cache_for_id(ID id)
     }
 
     rb_yjit_constant_state_changed(id);
+    rb_mjit_constant_state_changed(id);
 }
 
 static void