ZJIT: Remove PadPatchPoint instructions when lowering to LIR (#15974)

Basic blocks in LIR should only ever end in control flow instructions such as jump or return. PadPatchPoint is not control flow, so we should not emit it at the end of blocks when lowering.

4 files changed, 32 insertions, 5 deletions

diff --git a/zjit/src/backend/arm64/mod.rs b/zjit/src/backend/arm64/mod.rs
index ee15627d89..a1836ea9df 100644
--- a/zjit/src/backend/arm64/mod.rs
+++ b/zjit/src/backend/arm64/mod.rs
@@ -1752,7 +1752,7 @@ mod tests {
         asm.cret(val64);
 
         asm.frame_teardown(JIT_PRESERVED_REGS);
-        assert_disasm_snapshot!(lir_string(&mut asm), @r"
+        assert_disasm_snapshot!(lir_string(&mut asm), @"
         bb0:
           # bb0(): foo@/tmp/a.rb:1
           FrameSetup 1, x19, x21, x20
@@ -1765,6 +1765,7 @@ mod tests {
           Je bb0
           CRet v0
           FrameTeardown x19, x21, x20
+          PadPatchPoint
         ");
     }
 
diff --git a/zjit/src/backend/lir.rs b/zjit/src/backend/lir.rs
index f0fcece8a1..b2ec95a9d4 100644
--- a/zjit/src/backend/lir.rs
+++ b/zjit/src/backend/lir.rs
@@ -1745,11 +1745,25 @@ impl Assembler
         // Emit instructions with labels, expanding branch parameters
         let mut insns = Vec::with_capacity(ASSEMBLER_INSNS_CAPACITY);
 
-        for block in self.sorted_blocks() {
+        let blocks = self.sorted_blocks();
+        let num_blocks = blocks.len();
+
+        for (block_id, block) in blocks.iter().enumerate() {
+            // Entry blocks shouldn't ever be preceded by something that can
+            // stomp on this block.
+            if !block.is_entry {
+                insns.push(Insn::PadPatchPoint);
+            }
+
             // Process each instruction, expanding branch params if needed
             for insn in &block.insns {
                 self.expand_branch_insn(insn, &mut insns);
             }
+
+            // Make sure we don't stomp on the next function
+            if block_id == num_blocks - 1 {
+                insns.push(Insn::PadPatchPoint);
+            }
         }
         insns
     }
@@ -2211,6 +2225,11 @@ impl Assembler
         fn compile_exit(asm: &mut Assembler, exit: &SideExit) {
             let SideExit { pc, stack, locals } = exit;
 
+            // Side exit blocks are not part of the CFG at the moment,
+            // so we need to manually ensure that patchpoints get padded
+            // so that nobody stomps on us
+            asm.pad_patch_point();
+
             asm_comment!(asm, "save cfp->pc");
             asm.store(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_PC), *pc);
 
diff --git a/zjit/src/backend/x86_64/mod.rs b/zjit/src/backend/x86_64/mod.rs
index b045e0f3a3..d55dce1b9b 100644
--- a/zjit/src/backend/x86_64/mod.rs
+++ b/zjit/src/backend/x86_64/mod.rs
@@ -1173,7 +1173,7 @@ mod tests {
         asm.cret(val64);
 
         asm.frame_teardown(JIT_PRESERVED_REGS);
-        assert_disasm_snapshot!(lir_string(&mut asm), @r"
+        assert_disasm_snapshot!(lir_string(&mut asm), @"
         bb0:
           # bb0(): foo@/tmp/a.rb:1
           FrameSetup 1, r13, rbx, r12
@@ -1186,6 +1186,7 @@ mod tests {
           Je bb0
           CRet v0
           FrameTeardown r13, rbx, r12
+          PadPatchPoint
         ");
     }
 
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 41da154c1a..2038be808d 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -333,6 +333,8 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, version: IseqVersionRef, func
                     };
 
                     gen_if_false(&mut asm, val_opnd, branch_edge, fall_through_edge);
+                    assert!(asm.current_block().insns.last().unwrap().is_terminator());
+
                     asm.set_current_block(fall_through_target);
 
                     let label = jit.get_label(&mut asm, fall_through_target, block_id);
@@ -356,6 +358,8 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, version: IseqVersionRef, func
                     };
 
                     gen_if_true(&mut asm, val_opnd, branch_edge, fall_through_edge);
+                    assert!(asm.current_block().insns.last().unwrap().is_terminator());
+
                     asm.set_current_block(fall_through_target);
 
                     let label = jit.get_label(&mut asm, fall_through_target, block_id);
@@ -368,6 +372,8 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, version: IseqVersionRef, func
                         args: target.args.iter().map(|insn_id| jit.get_opnd(*insn_id)).collect()
                     };
                     gen_jump(&mut asm, branch_edge);
+                    assert!(asm.current_block().insns.last().unwrap().is_terminator());
+
                 },
                 _ => {
                     if let Err(last_snapshot) = gen_insn(cb, &mut jit, &mut asm, function, insn_id, &insn) {
@@ -382,8 +388,8 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, version: IseqVersionRef, func
                 }
             }
         }
-        // Make sure the last patch point has enough space to insert a jump
-        asm.pad_patch_point();
+        // Blocks should always end with control flow
+        assert!(asm.current_block().insns.last().unwrap().is_terminator());
     }
 
     // Generate code if everything can be compiled