YJIT: Assert no patch overlap on pos_marker (#9048)

7 files changed, 33 insertions, 17 deletions

diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index 7cd2449e73..52d844e121 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -1215,7 +1215,7 @@ impl Assembler
             // No bytes dropped, so the pos markers point to valid code
             for (insn_idx, pos) in pos_markers {
                 if let Insn::PosMarker(callback) = self.insns.get(insn_idx).unwrap() {
-                    callback(pos);
+                    callback(pos, &cb);
                 } else {
                     panic!("non-PosMarker in pos_markers insn_idx={insn_idx} {self:?}");
                 }
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index ce9c4f3f24..613e7048d4 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -327,7 +327,7 @@ impl From<CodePtr> for Target {
     }
 }
 
-type PosMarkerFn = Box<dyn Fn(CodePtr)>;
+type PosMarkerFn = Box<dyn Fn(CodePtr, &CodeBlock)>;
 
 /// YJIT IR instruction
 pub enum Insn {
@@ -1893,7 +1893,7 @@ impl Assembler {
     }
 
     //pub fn pos_marker<F: FnMut(CodePtr)>(&mut self, marker_fn: F)
-    pub fn pos_marker(&mut self, marker_fn: impl Fn(CodePtr) + 'static) {
+    pub fn pos_marker(&mut self, marker_fn: impl Fn(CodePtr, &CodeBlock) + 'static) {
         self.push_insn(Insn::PosMarker(Box::new(marker_fn)));
     }
 
diff --git a/yjit/src/backend/tests.rs b/yjit/src/backend/tests.rs
index 3278f33f63..01e87fe26c 100644
--- a/yjit/src/backend/tests.rs
+++ b/yjit/src/backend/tests.rs
@@ -318,7 +318,7 @@ fn test_no_pos_marker_callback_when_compile_fails() {
     let mut asm = Assembler::new();
 
     // Markers around code to exhaust memory limit
-    let fail_if_called = |_code_ptr| panic!("pos_marker callback should not be called");
+    let fail_if_called = |_code_ptr, _cb: &_| panic!("pos_marker callback should not be called");
     asm.pos_marker(fail_if_called);
     let zero = asm.load(0.into());
     let sum = asm.add(zero, 500.into());
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index 55bec66565..043139f492 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -827,7 +827,7 @@ impl Assembler
             // No bytes dropped, so the pos markers point to valid code
             for (insn_idx, pos) in pos_markers {
                 if let Insn::PosMarker(callback) = self.insns.get(insn_idx).unwrap() {
-                    callback(pos);
+                    callback(pos, &cb);
                 } else {
                     panic!("non-PosMarker in pos_markers insn_idx={insn_idx} {self:?}");
                 }
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 51a0ba3a31..41133e771f 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -245,13 +245,13 @@ impl JITState {
     fn perf_symbol_range_start(&self, asm: &mut Assembler, symbol_name: &str) {
         let symbol_name = symbol_name.to_string();
         let syms = self.perf_map.clone();
-        asm.pos_marker(move |start| syms.borrow_mut().push((start, None, symbol_name.clone())));
+        asm.pos_marker(move |start, _| syms.borrow_mut().push((start, None, symbol_name.clone())));
     }
 
     /// Mark the end address of a symbol to be reported to perf
     fn perf_symbol_range_end(&self, asm: &mut Assembler) {
         let syms = self.perf_map.clone();
-        asm.pos_marker(move |end| {
+        asm.pos_marker(move |end, _| {
             if let Some((_, ref mut end_store, _)) = syms.borrow_mut().last_mut() {
                 assert_eq!(None, *end_store);
                 *end_store = Some(end);
@@ -362,9 +362,13 @@ fn jit_prepare_routine_call(
 /// Record the current codeblock write position for rewriting into a jump into
 /// the outlined block later. Used to implement global code invalidation.
 fn record_global_inval_patch(asm: &mut Assembler, outline_block_target_pos: CodePtr) {
+    // We add a padding before pos_marker so that the previous patch will not overlap this.
+    // jump_to_next_insn() puts a patch point at the end of the block in fallthrough cases.
+    // In the fallthrough case, the next block should start with the same Context, so the
+    // patch is fine, but it should not overlap another patch.
     asm.pad_inval_patch();
-    asm.pos_marker(move |code_ptr| {
-        CodegenGlobals::push_global_inval_patch(code_ptr, outline_block_target_pos);
+    asm.pos_marker(move |code_ptr, cb| {
+        CodegenGlobals::push_global_inval_patch(code_ptr, outline_block_target_pos, cb);
     });
 }
 
@@ -8987,10 +8991,18 @@ impl CodegenGlobals {
         CodegenGlobals::get_instance().stub_exit_code
     }
 
-    pub fn push_global_inval_patch(i_pos: CodePtr, o_pos: CodePtr) {
+    pub fn push_global_inval_patch(inline_pos: CodePtr, outlined_pos: CodePtr, cb: &CodeBlock) {
+        if let Some(last_patch) = CodegenGlobals::get_instance().global_inval_patches.last() {
+            let patch_offset = inline_pos.as_offset() - last_patch.inline_patch_pos.as_offset();
+            assert!(
+                patch_offset < 0 || cb.jmp_ptr_bytes() as i64 <= patch_offset,
+                "patches should not overlap (patch_offset: {patch_offset})",
+            );
+        }
+
         let patch = CodepagePatch {
-            inline_patch_pos: i_pos,
-            outlined_target_pos: o_pos,
+            inline_patch_pos: inline_pos,
+            outlined_target_pos: outlined_pos,
         };
         CodegenGlobals::get_instance()
             .global_inval_patches
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 0ec5a0adca..0ab154d118 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -2842,7 +2842,7 @@ impl Assembler
         // so that we can move the closure below
         let entryref = entryref.clone();
 
-        self.pos_marker(move |code_ptr| {
+        self.pos_marker(move |code_ptr, _| {
             entryref.start_addr.set(Some(code_ptr));
         });
     }
@@ -2853,7 +2853,7 @@ impl Assembler
         // so that we can move the closure below
         let entryref = entryref.clone();
 
-        self.pos_marker(move |code_ptr| {
+        self.pos_marker(move |code_ptr, _| {
             entryref.end_addr.set(Some(code_ptr));
         });
     }
@@ -2865,7 +2865,7 @@ impl Assembler
         // so that we can move the closure below
         let branchref = branchref.clone();
 
-        self.pos_marker(move |code_ptr| {
+        self.pos_marker(move |code_ptr, _| {
             branchref.start_addr.set(Some(code_ptr));
         });
     }
@@ -2877,7 +2877,7 @@ impl Assembler
         // so that we can move the closure below
         let branchref = branchref.clone();
 
-        self.pos_marker(move |code_ptr| {
+        self.pos_marker(move |code_ptr, _| {
             branchref.end_addr.set(Some(code_ptr));
         });
     }
diff --git a/yjit/src/invariants.rs b/yjit/src/invariants.rs
index cffdfff117..13995ce5fb 100644
--- a/yjit/src/invariants.rs
+++ b/yjit/src/invariants.rs
@@ -524,7 +524,11 @@ pub extern "C" fn rb_yjit_tracing_invalidate_all() {
         patches.sort_by_cached_key(|patch| patch.inline_patch_pos.raw_ptr(cb));
         let mut last_patch_end = std::ptr::null();
         for patch in &patches {
-            assert!(last_patch_end <= patch.inline_patch_pos.raw_ptr(cb), "patches should not overlap");
+            let patch_pos = patch.inline_patch_pos.raw_ptr(cb);
+            assert!(
+                last_patch_end <= patch_pos,
+                "patches should not overlap (last_patch_end: {last_patch_end:?}, patch_pos: {patch_pos:?})",
+            );
 
             cb.set_write_ptr(patch.inline_patch_pos);
             cb.set_dropped_bytes(false);