ZJIT: Fix 30k if stmt test (#14446)

* ZJIT: Allow label generation above 19 bits

* Refactor emit_conditional_jump to use generate_branch
* Make branching functionality generic across Label and CodePtr

* ZJIT: Add > 19 bit jump test and helper function

* Remove an empty line

---------

Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>

2 files changed, 86 insertions, 60 deletions

diff --git a/zjit/src/asm/mod.rs b/zjit/src/asm/mod.rs
index 8efa5df270..e07ac0a48c 100644
--- a/zjit/src/asm/mod.rs
+++ b/zjit/src/asm/mod.rs
@@ -299,11 +299,15 @@ impl fmt::LowerHex for CodeBlock {
 impl CodeBlock {
     /// Stubbed CodeBlock for testing. Can't execute generated code.
     pub fn new_dummy() -> Self {
+        const DEFAULT_MEM_SIZE: usize = 1024;
+        CodeBlock::new_dummy_sized(DEFAULT_MEM_SIZE)
+    }
+
+    pub fn new_dummy_sized(mem_size: usize) -> Self {
         use std::ptr::NonNull;
         use crate::virtualmem::*;
         use crate::virtualmem::tests::TestingAllocator;
 
-        let mem_size = 1024;
         let alloc = TestingAllocator::new(mem_size);
         let mem_start: *const u8 = alloc.mem_start();
         let virt_mem = VirtualMem::new(alloc, 1, NonNull::new(mem_start as *mut u8).unwrap(), mem_size, 128 * 1024 * 1024);
@@ -388,4 +392,3 @@ mod tests
         assert_eq!(uimm_num_bits(u64::MAX), 64);
     }
 }
-
diff --git a/zjit/src/backend/arm64/mod.rs b/zjit/src/backend/arm64/mod.rs
index 9a23c2bd13..6e040e5214 100644
--- a/zjit/src/backend/arm64/mod.rs
+++ b/zjit/src/backend/arm64/mod.rs
@@ -706,70 +706,67 @@ impl Assembler
         /// Emit a conditional jump instruction to a specific target. This is
         /// called when lowering any of the conditional jump instructions.
         fn emit_conditional_jump<const CONDITION: u8>(cb: &mut CodeBlock, target: Target) {
+            fn generate_branch<const CONDITION: u8>(cb: &mut CodeBlock, src_addr: i64, dst_addr: i64) {
+                let num_insns = if bcond_offset_fits_bits((dst_addr - src_addr) / 4) {
+                    // If the jump offset fits into the conditional jump as
+                    // an immediate value and it's properly aligned, then we
+                    // can use the b.cond instruction directly. We're safe
+                    // to use as i32 here since we already checked that it
+                    // fits.
+                    let bytes = (dst_addr - src_addr) as i32;
+                    bcond(cb, CONDITION, InstructionOffset::from_bytes(bytes));
+
+                    // Here we're going to return 1 because we've only
+                    // written out 1 instruction.
+                    1
+                } else if b_offset_fits_bits((dst_addr - (src_addr + 4)) / 4) { // + 4 for bcond
+                    // If the jump offset fits into the unconditional jump as
+                    // an immediate value, we can use inverse b.cond + b.
+                    //
+                    // We're going to write out the inverse condition so
+                    // that if it doesn't match it will skip over the
+                    // instruction used for branching.
+                    bcond(cb, Condition::inverse(CONDITION), 2.into());
+                    b(cb, InstructionOffset::from_bytes((dst_addr - (src_addr + 4)) as i32)); // + 4 for bcond
+
+                    // We've only written out 2 instructions.
+                    2
+                } else {
+                    // Otherwise, we need to load the address into a
+                    // register and use the branch register instruction.
+                    let load_insns: i32 = emit_load_size(dst_addr as u64).into();
+
+                    // We're going to write out the inverse condition so
+                    // that if it doesn't match it will skip over the
+                    // instructions used for branching.
+                    bcond(cb, Condition::inverse(CONDITION), (load_insns + 2).into());
+                    emit_load_value(cb, Assembler::SCRATCH0, dst_addr as u64);
+                    br(cb, Assembler::SCRATCH0);
+
+                    // Here we'll return the number of instructions that it
+                    // took to write out the destination address + 1 for the
+                    // b.cond and 1 for the br.
+                    load_insns + 2
+                };
+
+                // We need to make sure we have at least 6 instructions for
+                // every kind of jump for invalidation purposes, so we're
+                // going to write out padding nop instructions here.
+                assert!(num_insns <= cb.conditional_jump_insns());
+                (num_insns..cb.conditional_jump_insns()).for_each(|_| nop(cb));
+            }
+
             match target {
                 Target::CodePtr(dst_ptr) => {
                     let dst_addr = dst_ptr.as_offset();
                     let src_addr = cb.get_write_ptr().as_offset();
-
-                    let num_insns = if bcond_offset_fits_bits((dst_addr - src_addr) / 4) {
-                        // If the jump offset fits into the conditional jump as
-                        // an immediate value and it's properly aligned, then we
-                        // can use the b.cond instruction directly. We're safe
-                        // to use as i32 here since we already checked that it
-                        // fits.
-                        let bytes = (dst_addr - src_addr) as i32;
-                        bcond(cb, CONDITION, InstructionOffset::from_bytes(bytes));
-
-                        // Here we're going to return 1 because we've only
-                        // written out 1 instruction.
-                        1
-                    } else if b_offset_fits_bits((dst_addr - (src_addr + 4)) / 4) { // + 4 for bcond
-                        // If the jump offset fits into the unconditional jump as
-                        // an immediate value, we can use inverse b.cond + b.
-                        //
-                        // We're going to write out the inverse condition so
-                        // that if it doesn't match it will skip over the
-                        // instruction used for branching.
-                        bcond(cb, Condition::inverse(CONDITION), 2.into());
-                        b(cb, InstructionOffset::from_bytes((dst_addr - (src_addr + 4)) as i32)); // + 4 for bcond
-
-                        // We've only written out 2 instructions.
-                        2
-                    } else {
-                        // Otherwise, we need to load the address into a
-                        // register and use the branch register instruction.
-                        let dst_addr = (dst_ptr.raw_ptr(cb) as usize).as_u64();
-                        let load_insns: i32 = emit_load_size(dst_addr).into();
-
-                        // We're going to write out the inverse condition so
-                        // that if it doesn't match it will skip over the
-                        // instructions used for branching.
-                        bcond(cb, Condition::inverse(CONDITION), (load_insns + 2).into());
-                        emit_load_value(cb, Assembler::SCRATCH0, dst_addr);
-                        br(cb, Assembler::SCRATCH0);
-
-                        // Here we'll return the number of instructions that it
-                        // took to write out the destination address + 1 for the
-                        // b.cond and 1 for the br.
-                        load_insns + 2
-                    };
-
-                    if let Target::CodePtr(_) = target {
-                        // We need to make sure we have at least 6 instructions for
-                        // every kind of jump for invalidation purposes, so we're
-                        // going to write out padding nop instructions here.
-                        assert!(num_insns <= cb.conditional_jump_insns());
-                        for _ in num_insns..cb.conditional_jump_insns() { nop(cb); }
-                    }
+                    generate_branch::<CONDITION>(cb, src_addr, dst_addr);
                 },
                 Target::Label(label_idx) => {
-                    // Here we're going to save enough space for ourselves and
-                    // then come back and write the instruction once we know the
-                    // offset. We're going to assume we can fit into a single
-                    // b.cond instruction. It will panic otherwise.
-                    cb.label_ref(label_idx, 4, |cb, src_addr, dst_addr| {
-                        let bytes: i32 = (dst_addr - (src_addr - 4)).try_into().unwrap();
-                        bcond(cb, CONDITION, InstructionOffset::from_bytes(bytes));
+                    // We save `cb.conditional_jump_insns` number of bytes since we may use up to that amount
+                    // `generate_branch` will pad the emitted branch instructions with `nop`s for each unused byte.
+                    cb.label_ref(label_idx, (cb.conditional_jump_insns() * 4) as usize, |cb, src_addr, dst_addr| {
+                        generate_branch::<CONDITION>(cb, src_addr - (cb.conditional_jump_insns() * 4) as i64, dst_addr);
                     });
                 },
                 Target::SideExit { .. } => {
@@ -2043,6 +2040,32 @@ mod tests {
     }
 
     #[test]
+    fn test_label_branch_generate_bounds() {
+        // The immediate in a conditional branch is a 19 bit unsigned integer
+        // which has a max value of 2^18 - 1.
+        const IMMEDIATE_MAX_VALUE: usize = 2usize.pow(18) - 1;
+
+        // `IMMEDIATE_MAX_VALUE` number of dummy instructions will be generated
+        // plus a compare, a jump instruction, and a label.
+        const MEMORY_REQUIRED: usize = (IMMEDIATE_MAX_VALUE + 8)*4;
+
+        let mut asm = Assembler::new();
+        let mut cb = CodeBlock::new_dummy_sized(MEMORY_REQUIRED);
+
+        let far_label = asm.new_label("far");
+
+        asm.cmp(Opnd::Reg(X0_REG), Opnd::UImm(1));
+        asm.je(far_label.clone());
+
+        (0..IMMEDIATE_MAX_VALUE).for_each(|_| {
+            asm.mov(Opnd::Reg(TEMP_REGS[0]), Opnd::Reg(TEMP_REGS[2]));
+        });
+
+        asm.write_label(far_label.clone());
+        asm.compile_with_num_regs(&mut cb, 1);
+    }
+
+    #[test]
     fn test_add_with_immediate() {
         let (mut asm, mut cb) = setup_asm();