TBZ and TBNZ for AArch64 (Shopify#434)

yjit/src/asm/arm64/inst/mod.rs

@@ -22,6 +22,7 @@ mod reg_pair;
 mod sbfm;
 mod shift_imm;
 mod sys_reg;
+mod test_bit;
 
 pub use atomic::Atomic;
 pub use branch::Branch;
@@ -44,3 +45,4 @@ pub use reg_pair::RegisterPair;
 pub use sbfm::SBFM;
 pub use shift_imm::ShiftImm;
 pub use sys_reg::SysReg;
+pub use test_bit::TestBit;

yjit/src/asm/arm64/inst/test_bit.rs

@@ -0,0 +1,135 @@
+/// The upper bit of the bit number to test.
+#[derive(Debug)]
+enum B5 {
+    /// When the bit number is below 32.
+    B532 = 0,
+
+    /// When the bit number is equal to or above 32.
+    B564 = 1
+}
+
+/// A convenience function so that we can convert the bit number directly into a
+/// B5 variant.
+impl From<u8> for B5 {
+    fn from(bit_num: u8) -> Self {
+        match bit_num {
+            0..=31 => B5::B532,
+            32..=63 => B5::B564,
+            _ => panic!("Invalid bit number: {}", bit_num)
+        }
+    }
+}
+
+/// The operation to perform for this instruction.
+enum Op {
+    /// The test bit zero operation.
+    TBZ = 0,
+
+    /// The test bit not zero operation.
+    TBNZ = 1
+}
+
+/// The struct that represents an A64 test bit instruction that can be encoded.
+///
+/// TBNZ/TBZ
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+/// | 31 30 29 28 | 27 26 25 24 | 23 22 21 20 | 19 18 17 16 | 15 14 13 12 | 11 10 09 08 | 07 06 05 04 | 03 02 01 00 |
+/// |     0  1  1    0  1  1                                                                                        |
+/// | b5                     op   b40............. imm14.......................................... rt.............. |
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+///
+pub struct TestBit {
+    /// The number of the register to test.
+    rt: u8,
+
+    /// The PC-relative offset to the target instruction in term of number of
+    /// instructions.
+    imm14: i16,
+
+    /// The lower 5 bits of the bit number to be tested.
+    b40: u8,
+
+    /// The operation to perform for this instruction.
+    op: Op,
+
+    /// The upper bit of the bit number to test.
+    b5: B5
+}
+
+impl TestBit {
+    /// TBNZ
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/TBNZ--Test-bit-and-Branch-if-Nonzero-?lang=en
+    pub fn tbnz(rt: u8, bit_num: u8, offset: i16) -> Self {
+        Self { rt, imm14: offset, b40: bit_num & 0b11111, op: Op::TBNZ, b5: bit_num.into() }
+    }
+
+    /// TBZ
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/TBZ--Test-bit-and-Branch-if-Zero-?lang=en
+    pub fn tbz(rt: u8, bit_num: u8, offset: i16) -> Self {
+        Self { rt, imm14: offset, b40: bit_num & 0b11111, op: Op::TBZ, b5: bit_num.into() }
+    }
+}
+
+/// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/Branches--Exception-Generating-and-System-instructions?lang=en
+const FAMILY: u32 = 0b11011;
+
+impl From<TestBit> for u32 {
+    /// Convert an instruction into a 32-bit value.
+    fn from(inst: TestBit) -> Self {
+        let b40 = (inst.b40 & 0b11111) as u32;
+        let mut imm14 = (inst.imm14 & ((1 << 13) - 1)) as u32;
+
+        if inst.imm14 < 0 {
+            imm14 |= (1 << 13);
+        }
+
+        0
+        | ((inst.b5 as u32) << 31)
+        | (FAMILY << 25)
+        | ((inst.op as u32) << 24)
+        | (b40 << 19)
+        | (imm14 << 5)
+        | inst.rt as u32
+    }
+}
+
+impl From<TestBit> for [u8; 4] {
+    /// Convert an instruction into a 4 byte array.
+    fn from(inst: TestBit) -> [u8; 4] {
+        let result: u32 = inst.into();
+        result.to_le_bytes()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_tbnz() {
+        let inst = TestBit::tbnz(0, 0, 0);
+        let result: u32 = inst.into();
+        assert_eq!(0x37000000, result);
+    }
+
+    #[test]
+    fn test_tbnz_negative() {
+        let inst = TestBit::tbnz(0, 0, -1);
+        let result: u32 = inst.into();
+        assert_eq!(0x3707ffe0, result);
+    }
+
+    #[test]
+    fn test_tbz() {
+        let inst = TestBit::tbz(0, 0, 0);
+        let result: u32 = inst.into();
+        assert_eq!(0x36000000, result);
+    }
+
+    #[test]
+    fn test_tbz_negative() {
+        let inst = TestBit::tbz(0, 0, -1);
+        let result: u32 = inst.into();
+        assert_eq!(0x3607ffe0, result);
+    }
+}

yjit/src/asm/arm64/mod.rs

@@ -934,6 +934,30 @@ pub fn ret(cb: &mut CodeBlock, rn: A64Opnd) {
     cb.write_bytes(&bytes);
 }
 
+/// TBNZ - test bit and branch if not zero
+pub fn tbnz(cb: &mut CodeBlock, rt: A64Opnd, bit_num: A64Opnd, offset: A64Opnd) {
+    let bytes: [u8; 4] = match (rt, bit_num, offset) {
+        (A64Opnd::Reg(rt), A64Opnd::UImm(bit_num), A64Opnd::Imm(offset)) => {
+            TestBit::tbnz(rt.reg_no, bit_num.try_into().unwrap(), offset.try_into().unwrap()).into()
+        },
+        _ => panic!("Invalid operand combination to tbnz instruction.")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
+/// TBZ - test bit and branch if zero
+pub fn tbz(cb: &mut CodeBlock, rt: A64Opnd, bit_num: A64Opnd, offset: A64Opnd) {
+    let bytes: [u8; 4] = match (rt, bit_num, offset) {
+        (A64Opnd::Reg(rt), A64Opnd::UImm(bit_num), A64Opnd::Imm(offset)) => {
+            TestBit::tbz(rt.reg_no, bit_num.try_into().unwrap(), offset.try_into().unwrap()).into()
+        },
+        _ => panic!("Invalid operand combination to tbz instruction.")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
 /// TST - test the bits of a register against a mask, then update flags
 pub fn tst(cb: &mut CodeBlock, rn: A64Opnd, rm: A64Opnd) {
     let bytes: [u8; 4] = match (rn, rm) {
@@ -1393,6 +1417,16 @@ mod tests {
         check_bytes("6a7d4093", |cb| sxtw(cb, X10, W11));
     }
 
+    #[test]
+    fn test_tbnz() {
+        check_bytes("4a005037", |cb| tbnz(cb, X10, A64Opnd::UImm(10), A64Opnd::Imm(2)));
+    }
+
+    #[test]
+    fn test_tbz() {
+        check_bytes("4a005036", |cb| tbz(cb, X10, A64Opnd::UImm(10), A64Opnd::Imm(2)));
+    }
+
     #[test]
     fn test_tst_register() {
         check_bytes("1f0001ea", |cb| tst(cb, X0, X1));