ZJIT: Support invalidating method redefinition

st0012 · st0012 · commit ad851484b7a7 · 2025-07-18T15:57:26.000+01:00
This commit adds support for the MethodRedefined invariant to be invalidated
when a method is redefined.

Changes:
- Added CME pointer to the MethodRedefined invariant in HIR
- Updated all places where MethodRedefined invariants are created to
    include the CME pointer
- Added handling for MethodRedefined invariants in gen_patch_point to
    call track_cme_assumption, which registers the patch point for
    invalidation when rb_zjit_cme_invalidate is called

This ensures that when a method is redefined, all JIT code that
depends on that method will be properly invalidated.
diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
@@ -1066,6 +1066,31 @@ def test
     }, call_threshold: 2
   end
 
+  # ZJIT currently only generates a MethodRedefined patch point when the method
+  # is called on the top-level self.
+  def test_method_redefinition_with_top_self
+    assert_runs '["original", "redefined"]', %q{
+      def foo
+        "original"
+      end
+
+      def test = foo
+
+      test; test
+
+      result1 = test
+
+      # Redefine the method
+      def foo
+        "redefined"
+      end
+
+      result2 = test
+
+      [result1, result2]
+    }, call_threshold: 2
+  end
+
   def test_module_name_with_guard_passes
     assert_compiles '"Integer"', %q{
       def test(mod)
diff --git a/vm_method.c b/vm_method.c
@@ -122,6 +122,7 @@ vm_cme_invalidate(rb_callable_method_entry_t *cme)
     RB_DEBUG_COUNTER_INC(cc_cme_invalidate);
 
     rb_yjit_cme_invalidate(cme);
+    rb_zjit_cme_invalidate(cme);
 }
 
 static int
diff --git a/zjit.h b/zjit.h
@@ -12,6 +12,7 @@ void rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, boo
 void rb_zjit_profile_insn(enum ruby_vminsn_type insn, rb_execution_context_t *ec);
 void rb_zjit_profile_enable(const rb_iseq_t *iseq);
 void rb_zjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop);
+void rb_zjit_cme_invalidate(const rb_callable_method_entry_t *cme);
 void rb_zjit_invalidate_ep_is_bp(const rb_iseq_t *iseq);
 void rb_zjit_iseq_mark(void *payload);
 void rb_zjit_iseq_update_references(void *payload);
@@ -21,6 +22,7 @@ static inline void rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_cont
 static inline void rb_zjit_profile_insn(enum ruby_vminsn_type insn, rb_execution_context_t *ec) {}
 static inline void rb_zjit_profile_enable(const rb_iseq_t *iseq) {}
 static inline void rb_zjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop) {}
+static inline void rb_zjit_cme_invalidate(const rb_callable_method_entry_t *cme) {}
 static inline void rb_zjit_invalidate_ep_is_bp(const rb_iseq_t *iseq) {}
 #endif // #if USE_YJIT
 
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
@@ -3,7 +3,7 @@ use std::rc::Rc;
 
 use crate::asm::Label;
 use crate::backend::current::{Reg, ALLOC_REGS};
-use crate::invariants::track_bop_assumption;
+use crate::invariants::{track_bop_assumption, track_cme_assumption};
 use crate::gc::{get_or_create_iseq_payload, append_gc_offsets};
 use crate::state::ZJITState;
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
@@ -494,6 +494,10 @@ fn gen_patch_point(jit: &mut JITState, asm: &mut Assembler, invariant: &Invarian
                 let side_exit_ptr = cb.resolve_label(label);
                 track_bop_assumption(klass, bop, code_ptr, side_exit_ptr);
             }
+            Invariant::MethodRedefined { klass: _, method: _, cme } => {
+                let side_exit_ptr = cb.resolve_label(label);
+                track_cme_assumption(cme, code_ptr, side_exit_ptr);
+            }
             _ => {
                 debug!("ZJIT: gen_patch_point: unimplemented invariant {invariant:?}");
                 return;
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
diff --git a/zjit/src/invariants.rs b/zjit/src/invariants.rs

Original file line number	Diff line number	Diff line change
`@@ -122,6 +122,7 @@ vm_cme_invalidate(rb_callable_method_entry_t *cme)`
`122`	`122`	`RB_DEBUG_COUNTER_INC(cc_cme_invalidate);`
`123`	`123`
`124`	`124`	`rb_yjit_cme_invalidate(cme);`
	`125`	`+ rb_zjit_cme_invalidate(cme);`
`125`	`126`	`}`
`126`	`127`
`127`	`128`	`static int`