TST, CMP, AND/ANDS with registers (https://github.com/Shopify/ruby/pull/301)

* Add TST instruction and AND/ANDS entrypoints for immediates

* TST/AND/ANDS for registers

* CMP instruction

5 files changed, 290 insertions, 0 deletions

diff --git a/yjit/src/asm/arm64/inst/data_imm.rs b/yjit/src/asm/arm64/inst/data_imm.rs
index 0d0a6ff325..950cf3421e 100644
--- a/yjit/src/asm/arm64/inst/data_imm.rs
+++ b/yjit/src/asm/arm64/inst/data_imm.rs
@@ -80,6 +80,12 @@ impl DataImm {
         }
     }
 
+    /// CMP (immediate)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/CMP--immediate---Compare--immediate---an-alias-of-SUBS--immediate--?lang=en
+    pub fn cmp(rn: u8, imm12: u16, num_bits: u8) -> Self {
+        Self::subs(31, rn, imm12, num_bits)
+    }
+
     /// SUB (immediate)
     /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/SUB--immediate---Subtract--immediate--?lang=en
     pub fn sub(rd: u8, rn: u8, imm12: u16, num_bits: u8) -> Self {
@@ -157,6 +163,13 @@ mod tests {
     }
 
     #[test]
+    fn test_cmp() {
+        let inst = DataImm::cmp(0, 7, 64);
+        let result: u32 = inst.into();
+        assert_eq!(0xf1001c1f, result);
+    }
+
+    #[test]
     fn test_sub() {
         let inst = DataImm::sub(0, 1, 7, 64);
         let result: u32 = inst.into();
diff --git a/yjit/src/asm/arm64/inst/data_reg.rs b/yjit/src/asm/arm64/inst/data_reg.rs
index 8635ab804b..40f026d1fd 100644
--- a/yjit/src/asm/arm64/inst/data_reg.rs
+++ b/yjit/src/asm/arm64/inst/data_reg.rs
@@ -86,6 +86,12 @@ impl DataReg {
         }
     }
 
+    /// CMP (shifted register)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/CMP--shifted-register---Compare--shifted-register---an-alias-of-SUBS--shifted-register--?lang=en
+    pub fn cmp(rn: u8, rm: u8, num_bits: u8) -> Self {
+        Self::subs(31, rn, rm, num_bits)
+    }
+
     /// SUB (shifted register)
     /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/SUB--shifted-register---Subtract--shifted-register--?lang=en
     pub fn sub(rd: u8, rn: u8, rm: u8, num_bits: u8) -> Self {
@@ -166,6 +172,13 @@ mod tests {
     }
 
     #[test]
+    fn test_cmp() {
+        let inst = DataReg::cmp(0, 1, 64);
+        let result: u32 = inst.into();
+        assert_eq!(0xeb01001f, result);
+    }
+
+    #[test]
     fn test_sub() {
         let inst = DataReg::sub(0, 1, 2, 64);
         let result: u32 = inst.into();
diff --git a/yjit/src/asm/arm64/inst/logical_imm.rs b/yjit/src/asm/arm64/inst/logical_imm.rs
index 63a4556d85..88de8ba4a1 100644
--- a/yjit/src/asm/arm64/inst/logical_imm.rs
+++ b/yjit/src/asm/arm64/inst/logical_imm.rs
@@ -49,6 +49,12 @@ impl LogicalImm {
     pub fn ands(rd: u8, rn: u8, imm: BitmaskImmediate, num_bits: u8) -> Self {
         Self { rd, rn, imm, opc: Opc::Ands, sf: num_bits.into() }
     }
+
+    /// TST (immediate)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/TST--immediate---Test-bits--immediate---an-alias-of-ANDS--immediate--?lang=en
+    pub fn tst(rn: u8, imm: BitmaskImmediate, num_bits: u8) -> Self {
+        Self::ands(31, rn, imm, num_bits)
+    }
 }
 
 /// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/Data-Processing----Immediate?lang=en#log_imm
@@ -94,4 +100,11 @@ mod tests {
         let result: u32 = inst.into();
         assert_eq!(0xf2400820, result);
     }
+
+    #[test]
+    fn test_tst() {
+        let inst = LogicalImm::tst(1, 7.try_into().unwrap(), 64);
+        let result: u32 = inst.into();
+        assert_eq!(0xf240083f, result);
+    }
 }
diff --git a/yjit/src/asm/arm64/inst/logical_reg.rs b/yjit/src/asm/arm64/inst/logical_reg.rs
new file mode 100644
index 0000000000..929d80b1a7
--- /dev/null
+++ b/yjit/src/asm/arm64/inst/logical_reg.rs
@@ -0,0 +1,125 @@
+use super::sf::Sf;
+
+/// The type of shift to perform on the second operand register.
+enum Shift {
+    LSL = 0b00, // logical shift left (unsigned)
+    LSR = 0b01, // logical shift right (unsigned)
+    ASR = 0b10, // arithmetic shift right (signed)
+    ROR = 0b11  // rotate right (unsigned)
+}
+
+// Which operation to perform.
+enum Opc {
+    /// The AND operation.
+    And = 0b00,
+
+    /// The ANDS operation.
+    Ands = 0b11
+}
+
+/// The struct that represents an A64 logical register instruction that can be
+/// encoded.
+///
+/// AND/ANDS (shifted register)
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+/// | 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  0  1  0          0                                                                          |
+/// | sf opc..                    shift    rm..............   imm6............... rn.............. rd.............. |
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+///
+pub struct LogicalReg {
+    /// The register number of the destination register.
+    rd: u8,
+
+    /// The register number of the first operand register.
+    rn: u8,
+
+    /// The amount to shift the second operand register.
+    imm6: u8,
+
+    /// The register number of the second operand register.
+    rm: u8,
+
+    /// The type of shift to perform on the second operand register.
+    shift: Shift,
+
+    /// The opcode for this instruction.
+    opc: Opc,
+
+    /// Whether or not this instruction is operating on 64-bit operands.
+    sf: Sf
+}
+
+impl LogicalReg {
+    /// AND (shifted register)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/AND--shifted-register---Bitwise-AND--shifted-register--?lang=en
+    pub fn and(rd: u8, rn: u8, rm: u8, num_bits: u8) -> Self {
+        Self { rd, rn, imm6: 0, rm, shift: Shift::LSL, opc: Opc::And, sf: num_bits.into() }
+    }
+
+    /// ANDS (shifted register)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/ANDS--shifted-register---Bitwise-AND--shifted-register---setting-flags-?lang=en
+    pub fn ands(rd: u8, rn: u8, rm: u8, num_bits: u8) -> Self {
+        Self { rd, rn, imm6: 0, rm, shift: Shift::LSL, opc: Opc::Ands, sf: num_bits.into() }
+    }
+
+    /// TST (shifted register)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/TST--shifted-register---Test--shifted-register---an-alias-of-ANDS--shifted-register--?lang=en
+    pub fn tst(rn: u8, rm: u8, num_bits: u8) -> Self {
+        Self { rd: 31, rn, imm6: 0, rm, shift: Shift::LSL, opc: Opc::Ands, sf: num_bits.into() }
+    }
+}
+
+/// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/Data-Processing----Register?lang=en
+const FAMILY: u32 = 0b0101;
+
+impl From<LogicalReg> for u32 {
+    /// Convert an instruction into a 32-bit value.
+    fn from(inst: LogicalReg) -> Self {
+        let imm6 = (inst.imm6 as u32) & ((1 << 6) - 1);
+
+        0
+        | ((inst.sf as u32) << 31)
+        | ((inst.opc as u32) << 29)
+        | (FAMILY << 25)
+        | ((inst.shift as u32) << 22)
+        | ((inst.rm as u32) << 16)
+        | (imm6 << 10)
+        | ((inst.rn as u32) << 5)
+        | inst.rd as u32
+    }
+}
+
+impl From<LogicalReg> for [u8; 4] {
+    /// Convert an instruction into a 4 byte array.
+    fn from(inst: LogicalReg) -> [u8; 4] {
+        let result: u32 = inst.into();
+        result.to_le_bytes()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_and() {
+        let inst = LogicalReg::and(0, 1, 2, 64);
+        let result: u32 = inst.into();
+        assert_eq!(0x8a020020, result);
+    }
+
+    #[test]
+    fn test_ands() {
+        let inst = LogicalReg::ands(0, 1, 2, 64);
+        let result: u32 = inst.into();
+        assert_eq!(0xea020020, result);
+    }
+
+    #[test]
+    fn test_tst() {
+        let inst = LogicalReg::tst(0, 1, 64);
+        let result: u32 = inst.into();
+        assert_eq!(0xea01001f, result);
+    }
+}
diff --git a/yjit/src/asm/arm64/inst/mod.rs b/yjit/src/asm/arm64/inst/mod.rs
index 83cdd26d1d..7d05f28604 100644
--- a/yjit/src/asm/arm64/inst/mod.rs
+++ b/yjit/src/asm/arm64/inst/mod.rs
@@ -6,6 +6,7 @@ mod data_imm;
 mod data_reg;
 mod load;
 mod logical_imm;
+mod logical_reg;
 mod mov;
 mod sf;
 mod store;
@@ -18,6 +19,8 @@ use call::Call;
 use data_imm::DataImm;
 use data_reg::DataReg;
 use load::Load;
+use logical_imm::LogicalImm;
+use logical_reg::LogicalReg;
 use mov::Mov;
 use store::Store;
 
@@ -85,6 +88,50 @@ pub fn adds(cb: &mut CodeBlock, rd: A64Opnd, rn: A64Opnd, rm: A64Opnd) {
     cb.write_bytes(&bytes);
 }
 
+/// AND - and rn and rm, put the result in rd, don't update flags
+pub fn and(cb: &mut CodeBlock, rd: A64Opnd, rn: A64Opnd, rm: A64Opnd) {
+    let bytes: [u8; 4] = match (rd, rn, rm) {
+        (A64Opnd::Reg(rd), A64Opnd::Reg(rn), A64Opnd::Reg(rm)) => {
+            assert!(
+                rd.num_bits == rn.num_bits && rn.num_bits == rm.num_bits,
+                "All operands must be of the same size."
+            );
+
+            LogicalReg::and(rd.reg_no, rn.reg_no, rm.reg_no, rd.num_bits).into()
+        },
+        (A64Opnd::Reg(rd), A64Opnd::Reg(rn), A64Opnd::UImm(imm)) => {
+            assert!(rd.num_bits == rn.num_bits, "rd and rn must be of the same size.");
+
+            LogicalImm::and(rd.reg_no, rn.reg_no, imm.try_into().unwrap(), rd.num_bits).into()
+        },
+        _ => panic!("Invalid operand combination to and instruction."),
+    };
+
+    cb.write_bytes(&bytes);
+}
+
+/// ANDS - and rn and rm, put the result in rd, update flags
+pub fn ands(cb: &mut CodeBlock, rd: A64Opnd, rn: A64Opnd, rm: A64Opnd) {
+    let bytes: [u8; 4] = match (rd, rn, rm) {
+        (A64Opnd::Reg(rd), A64Opnd::Reg(rn), A64Opnd::Reg(rm)) => {
+            assert!(
+                rd.num_bits == rn.num_bits && rn.num_bits == rm.num_bits,
+                "All operands must be of the same size."
+            );
+
+            LogicalReg::ands(rd.reg_no, rn.reg_no, rm.reg_no, rd.num_bits).into()
+        },
+        (A64Opnd::Reg(rd), A64Opnd::Reg(rn), A64Opnd::UImm(imm)) => {
+            assert!(rd.num_bits == rn.num_bits, "rd and rn must be of the same size.");
+
+            LogicalImm::ands(rd.reg_no, rn.reg_no, imm.try_into().unwrap(), rd.num_bits).into()
+        },
+        _ => panic!("Invalid operand combination to ands instruction."),
+    };
+
+    cb.write_bytes(&bytes);
+}
+
 /// BL - branch with link (offset is number of instructions to jump)
 pub fn bl(cb: &mut CodeBlock, imm26: A64Opnd) {
     let bytes: [u8; 4] = match imm26 {
@@ -109,6 +156,28 @@ pub fn br(cb: &mut CodeBlock, rn: A64Opnd) {
     cb.write_bytes(&bytes);
 }
 
+/// CMP - compare rn and rm, update flags
+pub fn cmp(cb: &mut CodeBlock, rn: A64Opnd, rm: A64Opnd) {
+    let bytes: [u8; 4] = match (rn, rm) {
+        (A64Opnd::Reg(rn), A64Opnd::Reg(rm)) => {
+            assert!(
+                rn.num_bits == rm.num_bits,
+                "All operands must be of the same size."
+            );
+
+            DataReg::cmp(rn.reg_no, rm.reg_no, rn.num_bits).into()
+        },
+        (A64Opnd::Reg(rn), A64Opnd::UImm(imm12)) => {
+            assert!(uimm_fits_bits(imm12, 12), "The immediate operand must be 12 bits or less.");
+
+            DataImm::cmp(rn.reg_no, imm12 as u16, rn.num_bits).into()
+        },
+        _ => panic!("Invalid operand combination to cmp instruction."),
+    };
+
+    cb.write_bytes(&bytes);
+}
+
 /// LDADDAL - atomic add with acquire and release semantics
 pub fn ldaddal(cb: &mut CodeBlock, rs: A64Opnd, rt: A64Opnd, rn: A64Opnd) {
     let bytes: [u8; 4] = match (rs, rt, rn) {
@@ -241,6 +310,23 @@ pub fn ret(cb: &mut CodeBlock, rn: A64Opnd) {
     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) {
+        (A64Opnd::Reg(rn), A64Opnd::Reg(rm)) => {
+            assert!(rn.num_bits == rm.num_bits, "All operands must be of the same size.");
+
+            LogicalReg::tst(rn.reg_no, rm.reg_no, rn.num_bits).into()
+        },
+        (A64Opnd::Reg(rn), A64Opnd::UImm(imm)) => {
+            LogicalImm::tst(rn.reg_no, imm.try_into().unwrap(), rn.num_bits).into()
+        },
+        _ => panic!("Invalid operand combination to tst instruction."),
+    };
+
+    cb.write_bytes(&bytes);
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -296,6 +382,26 @@ mod tests {
     }
 
     #[test]
+    fn test_and_register() {
+        check_bytes("2000028a", |cb| and(cb, X0, X1, X2));
+    }
+
+    #[test]
+    fn test_and_immediate() {
+        check_bytes("20084092", |cb| and(cb, X0, X1, A64Opnd::new_uimm(7)));
+    }
+
+    #[test]
+    fn test_ands_register() {
+        check_bytes("200002ea", |cb| ands(cb, X0, X1, X2));
+    }
+
+    #[test]
+    fn test_ands_immediate() {
+        check_bytes("200840f2", |cb| ands(cb, X0, X1, A64Opnd::new_uimm(7)));
+    }
+
+    #[test]
     fn test_bl() {
         check_bytes("00040094", |cb| bl(cb, A64Opnd::new_imm(1024)));
     }
@@ -306,6 +412,16 @@ mod tests {
     }
 
     #[test]
+    fn test_cmp_register() {
+        check_bytes("5f010beb", |cb| cmp(cb, X10, X11));
+    }
+
+    #[test]
+    fn test_cmp_immediate() {
+        check_bytes("5f3900f1", |cb| cmp(cb, X10, A64Opnd::new_uimm(14)));
+    }
+
+    #[test]
     fn test_ldaddal() {
         check_bytes("8b01eaf8", |cb| ldaddal(cb, X10, X11, X12));
     }
@@ -359,4 +475,14 @@ mod tests {
     fn test_subs_immediate() {
         check_bytes("201c00f1", |cb| subs(cb, X0, X1, A64Opnd::new_uimm(7)));
     }
+
+    #[test]
+    fn test_tst_register() {
+        check_bytes("1f0001ea", |cb| tst(cb, X0, X1));
+    }
+
+    #[test]
+    fn test_tst_immediate() {
+        check_bytes("3f0840f2", |cb| tst(cb, X1, A64Opnd::new_uimm(7)));
+    }
 }