BLR instruction for AArch64 (https://github.com/Shopify/ruby/pull/325)

2 files changed, 34 insertions, 4 deletions

diff --git a/yjit/src/asm/arm64/inst/branch.rs b/yjit/src/asm/arm64/inst/branch.rs
index 7f93f5e201..f15ef2a9b0 100644
--- a/yjit/src/asm/arm64/inst/branch.rs
+++ b/yjit/src/asm/arm64/inst/branch.rs
@@ -1,10 +1,13 @@
 /// Which operation to perform.
 enum Op {
     /// Perform a BR instruction.
-    Br = 0b00,
+    BR = 0b00,
+
+    /// Perform a BLR instruction.
+    BLR = 0b01,
 
     /// Perform a RET instruction.
-    Ret = 0b10
+    RET = 0b10
 }
 
 /// The struct that represents an A64 branch instruction that can be encoded.
@@ -27,13 +30,19 @@ impl Branch {
     /// BR
     /// https://developer.arm.com/documentation/ddi0602/2022-03/Base-Instructions/BR--Branch-to-Register-?lang=en
     pub fn br(rn: u8) -> Self {
-        Self { rn, op: Op::Br }
+        Self { rn, op: Op::BR }
+    }
+
+    /// BLR
+    /// https://developer.arm.com/documentation/ddi0602/2022-03/Base-Instructions/BLR--Branch-with-Link-to-Register-?lang=en
+    pub fn blr(rn: u8) -> Self {
+        Self { rn, op: Op::BLR }
     }
 
     /// RET
     /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/RET--Return-from-subroutine-?lang=en
     pub fn ret(rn: u8) -> Self {
-        Self { rn, op: Op::Ret }
+        Self { rn, op: Op::RET }
     }
 }
 
@@ -72,6 +81,12 @@ mod tests {
     }
 
     #[test]
+    fn test_blr() {
+        let result: u32 = Branch::blr(0).into();
+        assert_eq!(0xd63f0000, result);
+    }
+
+    #[test]
     fn test_ret() {
         let result: u32 = Branch::ret(30).into();
         assert_eq!(0xd65f03C0, result);
diff --git a/yjit/src/asm/arm64/mod.rs b/yjit/src/asm/arm64/mod.rs
index 333fce495f..8be7e6f568 100644
--- a/yjit/src/asm/arm64/mod.rs
+++ b/yjit/src/asm/arm64/mod.rs
@@ -226,6 +226,16 @@ pub fn bl(cb: &mut CodeBlock, imm26: A64Opnd) {
     cb.write_bytes(&bytes);
 }
 
+/// BLR - branch with link to a register
+pub fn blr(cb: &mut CodeBlock, rn: A64Opnd) {
+    let bytes: [u8; 4] = match rn {
+        A64Opnd::Reg(rn) => Branch::blr(rn.reg_no).into(),
+        _ => panic!("Invalid operand to blr instruction."),
+    };
+
+    cb.write_bytes(&bytes);
+}
+
 /// BR - branch to a register
 pub fn br(cb: &mut CodeBlock, rn: A64Opnd) {
     let bytes: [u8; 4] = match rn {
@@ -843,6 +853,11 @@ mod tests {
     }
 
     #[test]
+    fn test_blr() {
+        check_bytes("80023fd6", |cb| blr(cb, X20));
+    }
+
+    #[test]
     fn test_br() {
         check_bytes("80021fd6", |cb| br(cb, X20));
     }