Better splitting for Op::Test on AArch64 (https://github.com/Shopify/ruby/pull/335)

2 files changed, 101 insertions, 3 deletions

diff --git a/yjit/src/asm/arm64/arg/bitmask_imm.rs b/yjit/src/asm/arm64/arg/bitmask_imm.rs
index 7e5a21c7b4..847b735eaa 100644
--- a/yjit/src/asm/arm64/arg/bitmask_imm.rs
+++ b/yjit/src/asm/arm64/arg/bitmask_imm.rs
@@ -41,13 +41,19 @@ impl TryFrom<u64> for BitmaskImmediate {
 
     /// Attempt to convert a u64 into a BitmaskImm.
     fn try_from(value: u64) -> Result<Self, Self::Error> {
+        // 0 is not encodable as a bitmask immediate. Immediately return here so
+        // that we don't have any issues with underflow.
+        if value == 0 {
+            return Err(());
+        }
+
         /// Is this number's binary representation all 1s?
         fn is_mask(imm: u64) -> bool {
             if imm == u64::MAX { true } else { ((imm + 1) & imm) == 0 }
         }
 
-        /// Is this number's binary representation one or more 1s followed by one or
-        /// more 0s?
+        /// Is this number's binary representation one or more 1s followed by
+        /// one or more 0s?
         fn is_shifted_mask(imm: u64) -> bool {
             is_mask((imm - 1) | imm)
         }
diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index 1f93441c03..b1f4d63d0f 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -249,7 +249,33 @@ impl Assembler
                         _ => asm.load(opnds[0])
                     };
 
-                    asm.test(opnd0, opnds[1]);
+                    // The second value must be either a register or an
+                    // unsigned immediate that can be encoded as a bitmask
+                    // immediate. If it's not one of those, we'll need to load
+                    // it first.
+                    let opnd1 = match opnds[1] {
+                        Opnd::Reg(_) | Opnd::InsnOut { .. } => opnds[1],
+                        Opnd::Mem(_) => asm.load(opnds[1]),
+                        Opnd::Imm(imm) => {
+                            if imm <= 0 {
+                                asm.load(opnds[1])
+                            } else if BitmaskImmediate::try_from(imm as u64).is_ok() {
+                                Opnd::UImm(imm as u64)
+                            } else {
+                                asm.load(opnds[1])
+                            }
+                        },
+                        Opnd::UImm(uimm) => {
+                            if BitmaskImmediate::try_from(uimm).is_ok() {
+                                opnds[1]
+                            } else {
+                                asm.load(opnds[1])
+                            }
+                        },
+                        Opnd::None | Opnd::Value(_) => unreachable!()
+                    };
+
+                    asm.test(opnd0, opnd1);
                 },
                 _ => {
                     asm.push_insn(op, opnds, target, text, pos_marker);
@@ -790,6 +816,72 @@ mod tests {
     }
 
     #[test]
+    fn test_emit_test() {
+        let (mut asm, mut cb) = setup_asm();
+
+        asm.test(Opnd::Reg(X0_REG), Opnd::Reg(X1_REG));
+        asm.compile_with_num_regs(&mut cb, 0);
+
+        // Assert that only one instruction was written.
+        assert_eq!(4, cb.get_write_pos());
+    }
+
+    #[test]
+    fn test_emit_test_with_encodable_unsigned_immediate() {
+        let (mut asm, mut cb) = setup_asm();
+
+        asm.test(Opnd::Reg(X0_REG), Opnd::UImm(7));
+        asm.compile_with_num_regs(&mut cb, 0);
+
+        // Assert that only one instruction was written.
+        assert_eq!(4, cb.get_write_pos());
+    }
+
+    #[test]
+    fn test_emit_test_with_unencodable_unsigned_immediate() {
+        let (mut asm, mut cb) = setup_asm();
+
+        asm.test(Opnd::Reg(X0_REG), Opnd::UImm(5));
+        asm.compile_with_num_regs(&mut cb, 1);
+
+        // Assert that a load and a test instruction were written.
+        assert_eq!(8, cb.get_write_pos());
+    }
+
+    #[test]
+    fn test_emit_test_with_encodable_signed_immediate() {
+        let (mut asm, mut cb) = setup_asm();
+
+        asm.test(Opnd::Reg(X0_REG), Opnd::Imm(7));
+        asm.compile_with_num_regs(&mut cb, 0);
+
+        // Assert that only one instruction was written.
+        assert_eq!(4, cb.get_write_pos());
+    }
+
+    #[test]
+    fn test_emit_test_with_unencodable_signed_immediate() {
+        let (mut asm, mut cb) = setup_asm();
+
+        asm.test(Opnd::Reg(X0_REG), Opnd::Imm(5));
+        asm.compile_with_num_regs(&mut cb, 1);
+
+        // Assert that a load and a test instruction were written.
+        assert_eq!(8, cb.get_write_pos());
+    }
+
+    #[test]
+    fn test_emit_test_with_negative_signed_immediate() {
+        let (mut asm, mut cb) = setup_asm();
+
+        asm.test(Opnd::Reg(X0_REG), Opnd::Imm(-7));
+        asm.compile_with_num_regs(&mut cb, 1);
+
+        // Assert that a load and a test instruction were written.
+        assert_eq!(8, cb.get_write_pos());
+    }
+
+    #[test]
     #[cfg(feature = "disasm")]
     fn test_simple_disasm() -> std::result::Result<(), capstone::Error> {
         // Test drive Capstone with simple input