LDRH and STRH for AArch64 (https://github.com/Shopify/ruby/pull/438)

3 files changed, 298 insertions, 0 deletions

diff --git a/yjit/src/asm/arm64/inst/halfword_imm.rs b/yjit/src/asm/arm64/inst/halfword_imm.rs
new file mode 100644
index 0000000000..675e33d4a8
--- /dev/null
+++ b/yjit/src/asm/arm64/inst/halfword_imm.rs
@@ -0,0 +1,176 @@
+/// Whether this is a load or a store.
+enum Op {
+    Load = 1,
+    Store = 0
+}
+
+/// The type of indexing to perform for this instruction.
+enum Index {
+    /// No indexing.
+    None = 0b00,
+
+    /// Mutate the register after the read.
+    PostIndex = 0b01,
+
+    /// Mutate the register before the read.
+    PreIndex = 0b11
+}
+
+/// The struct that represents an A64 halfword instruction that can be encoded.
+///
+/// LDRH/STRH
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+/// | 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  1    1  0  0  1    0                                                                                |
+/// |                                op imm12.................................... rn.............. rt.............. |
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+///
+/// LDRH (pre-index/post-index)
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+/// | 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  1    1  0  0  0    0     0                                                                          |
+/// |                                op    imm9..........................   index rn.............. rt.............. |
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+///
+pub struct HalfwordImm {
+    /// The number of the 32-bit register to be loaded.
+    rt: u8,
+
+    /// The number of the 64-bit base register to calculate the memory address.
+    rn: u8,
+
+    /// The type of indexing to perform for this instruction.
+    index: Index,
+
+    /// The immediate offset from the base register.
+    imm: i16,
+
+    /// The operation to perform.
+    op: Op
+}
+
+impl HalfwordImm {
+    /// LDRH
+    /// https://developer.arm.com/documentation/ddi0602/2022-06/Base-Instructions/LDRH--immediate---Load-Register-Halfword--immediate--
+    pub fn ldrh(rt: u8, rn: u8, imm12: i16) -> Self {
+        Self { rt, rn, index: Index::None, imm: imm12, op: Op::Load }
+    }
+
+    /// LDRH (pre-index)
+    /// https://developer.arm.com/documentation/ddi0602/2022-06/Base-Instructions/LDRH--immediate---Load-Register-Halfword--immediate--
+    pub fn ldrh_pre(rt: u8, rn: u8, imm9: i16) -> Self {
+        Self { rt, rn, index: Index::PreIndex, imm: imm9, op: Op::Load }
+    }
+
+    /// LDRH (post-index)
+    /// https://developer.arm.com/documentation/ddi0602/2022-06/Base-Instructions/LDRH--immediate---Load-Register-Halfword--immediate--
+    pub fn ldrh_post(rt: u8, rn: u8, imm9: i16) -> Self {
+        Self { rt, rn, index: Index::PostIndex, imm: imm9, op: Op::Load }
+    }
+
+    /// STRH
+    /// https://developer.arm.com/documentation/ddi0602/2022-06/Base-Instructions/STRH--immediate---Store-Register-Halfword--immediate--
+    pub fn strh(rt: u8, rn: u8, imm12: i16) -> Self {
+        Self { rt, rn, index: Index::None, imm: imm12, op: Op::Store }
+    }
+
+    /// STRH (pre-index)
+    /// https://developer.arm.com/documentation/ddi0602/2022-06/Base-Instructions/STRH--immediate---Store-Register-Halfword--immediate--
+    pub fn strh_pre(rt: u8, rn: u8, imm9: i16) -> Self {
+        Self { rt, rn, index: Index::PreIndex, imm: imm9, op: Op::Store }
+    }
+
+    /// STRH (post-index)
+    /// https://developer.arm.com/documentation/ddi0602/2022-06/Base-Instructions/STRH--immediate---Store-Register-Halfword--immediate--
+    pub fn strh_post(rt: u8, rn: u8, imm9: i16) -> Self {
+        Self { rt, rn, index: Index::PostIndex, imm: imm9, op: Op::Store }
+    }
+}
+
+/// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/Loads-and-Stores?lang=en
+const FAMILY: u32 = 0b111100;
+
+impl From<HalfwordImm> for u32 {
+    /// Convert an instruction into a 32-bit value.
+    fn from(inst: HalfwordImm) -> Self {
+        let (mut opc, imm) = match inst.index {
+            Index::None => {
+                let mut imm12 = ((inst.imm / 2) as u32) & ((1 << 12) - 1);
+                (0b100, imm12)
+            },
+            Index::PreIndex | Index::PostIndex => {
+                let mut imm9 = (inst.imm as u32) & ((1 << 9) - 1);
+                (0b000, (imm9 << 2) | (inst.index as u32))
+            }
+        };
+
+        0
+        | (FAMILY << 25)
+        | ((opc | (inst.op as u32)) << 22)
+        | (imm << 10)
+        | ((inst.rn as u32) << 5)
+        | (inst.rt as u32)
+    }
+}
+
+impl From<HalfwordImm> for [u8; 4] {
+    /// Convert an instruction into a 4 byte array.
+    fn from(inst: HalfwordImm) -> [u8; 4] {
+        let result: u32 = inst.into();
+        result.to_le_bytes()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_ldrh() {
+        let inst = HalfwordImm::ldrh(0, 1, 8);
+        let result: u32 = inst.into();
+        assert_eq!(0x79401020, result);
+    }
+
+    #[test]
+    fn test_ldrh_pre() {
+        let inst = HalfwordImm::ldrh_pre(0, 1, 16);
+        let result: u32 = inst.into();
+        assert_eq!(0x78410c20, result);
+    }
+
+    #[test]
+    fn test_ldrh_post() {
+        let inst = HalfwordImm::ldrh_post(0, 1, 24);
+        let result: u32 = inst.into();
+        assert_eq!(0x78418420, result);
+    }
+
+    #[test]
+    fn test_ldrh_post_negative() {
+        let inst = HalfwordImm::ldrh_post(0, 1, -24);
+        let result: u32 = inst.into();
+        assert_eq!(0x785e8420, result);
+    }
+
+    #[test]
+    fn test_strh() {
+        let inst = HalfwordImm::strh(0, 1, 0);
+        let result: u32 = inst.into();
+        assert_eq!(0x79000020, result);
+    }
+
+    #[test]
+    fn test_strh_pre() {
+        let inst = HalfwordImm::strh_pre(0, 1, 0);
+        let result: u32 = inst.into();
+        assert_eq!(0x78000c20, result);
+    }
+
+    #[test]
+    fn test_strh_post() {
+        let inst = HalfwordImm::strh_post(0, 1, 0);
+        let result: u32 = inst.into();
+        assert_eq!(0x78000420, result);
+    }
+}
diff --git a/yjit/src/asm/arm64/inst/mod.rs b/yjit/src/asm/arm64/inst/mod.rs
index ab41464013..f4c27a5102 100644
--- a/yjit/src/asm/arm64/inst/mod.rs
+++ b/yjit/src/asm/arm64/inst/mod.rs
@@ -9,6 +9,7 @@ mod call;
 mod conditional;
 mod data_imm;
 mod data_reg;
+mod halfword_imm;
 mod load_literal;
 mod load_register;
 mod load_store;
@@ -30,6 +31,7 @@ pub use call::Call;
 pub use conditional::Conditional;
 pub use data_imm::DataImm;
 pub use data_reg::DataReg;
+pub use halfword_imm::HalfwordImm;
 pub use load_literal::LoadLiteral;
 pub use load_register::LoadRegister;
 pub use load_store::LoadStore;
diff --git a/yjit/src/asm/arm64/mod.rs b/yjit/src/asm/arm64/mod.rs
index fb07498ce2..cf898d2b5a 100644
--- a/yjit/src/asm/arm64/mod.rs
+++ b/yjit/src/asm/arm64/mod.rs
@@ -423,6 +423,51 @@ pub fn ldr_literal(cb: &mut CodeBlock, rt: A64Opnd, rn: i32) {
     cb.write_bytes(&bytes);
 }
 
+/// LDRH - load a halfword from memory
+pub fn ldrh(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
+    let bytes: [u8; 4] = match (rt, rn) {
+        (A64Opnd::Reg(rt), A64Opnd::Mem(rn)) => {
+            assert_eq!(rt.num_bits, 32, "Expected to be loading a halfword");
+            assert!(imm_fits_bits(rn.disp.into(), 12), "The displacement must be 12 bits or less.");
+
+            HalfwordImm::ldrh(rt.reg_no, rn.base_reg_no, rn.disp as i16).into()
+        },
+        _ => panic!("Invalid operand combination to ldrh instruction.")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
+/// LDRH (pre-index) - load a halfword from memory, update the base pointer before loading it
+pub fn ldrh_pre(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
+    let bytes: [u8; 4] = match (rt, rn) {
+        (A64Opnd::Reg(rt), A64Opnd::Mem(rn)) => {
+            assert_eq!(rt.num_bits, 32, "Expected to be loading a halfword");
+            assert!(imm_fits_bits(rn.disp.into(), 9), "The displacement must be 9 bits or less.");
+
+            HalfwordImm::ldrh_pre(rt.reg_no, rn.base_reg_no, rn.disp as i16).into()
+        },
+        _ => panic!("Invalid operand combination to ldrh instruction.")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
+/// LDRH (post-index) - load a halfword from memory, update the base pointer after loading it
+pub fn ldrh_post(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
+    let bytes: [u8; 4] = match (rt, rn) {
+        (A64Opnd::Reg(rt), A64Opnd::Mem(rn)) => {
+            assert_eq!(rt.num_bits, 32, "Expected to be loading a halfword");
+            assert!(imm_fits_bits(rn.disp.into(), 9), "The displacement must be 9 bits or less.");
+
+            HalfwordImm::ldrh_post(rt.reg_no, rn.base_reg_no, rn.disp as i16).into()
+        },
+        _ => panic!("Invalid operand combination to ldrh instruction.")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
 /// Whether or not a memory address displacement fits into the maximum number of
 /// bits such that it can be used without loading it into a register first.
 pub fn mem_disp_fits_bits(disp: i32) -> bool {
@@ -741,6 +786,51 @@ pub fn str_pre(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
     cb.write_bytes(&bytes);
 }
 
+/// STRH - store a halfword into memory
+pub fn strh(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
+    let bytes: [u8; 4] = match (rt, rn) {
+        (A64Opnd::Reg(rt), A64Opnd::Mem(rn)) => {
+            assert_eq!(rt.num_bits, 32, "Expected to be loading a halfword");
+            assert!(imm_fits_bits(rn.disp.into(), 12), "The displacement must be 12 bits or less.");
+
+            HalfwordImm::strh(rt.reg_no, rn.base_reg_no, rn.disp as i16).into()
+        },
+        _ => panic!("Invalid operand combination to strh instruction.")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
+/// STRH (pre-index) - store a halfword into memory, update the base pointer before loading it
+pub fn strh_pre(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
+    let bytes: [u8; 4] = match (rt, rn) {
+        (A64Opnd::Reg(rt), A64Opnd::Mem(rn)) => {
+            assert_eq!(rt.num_bits, 32, "Expected to be loading a halfword");
+            assert!(imm_fits_bits(rn.disp.into(), 9), "The displacement must be 9 bits or less.");
+
+            HalfwordImm::strh_pre(rt.reg_no, rn.base_reg_no, rn.disp as i16).into()
+        },
+        _ => panic!("Invalid operand combination to strh instruction.")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
+/// STRH (post-index) - store a halfword into memory, update the base pointer after loading it
+pub fn strh_post(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
+    let bytes: [u8; 4] = match (rt, rn) {
+        (A64Opnd::Reg(rt), A64Opnd::Mem(rn)) => {
+            assert_eq!(rt.num_bits, 32, "Expected to be loading a halfword");
+            assert!(imm_fits_bits(rn.disp.into(), 9), "The displacement must be 9 bits or less.");
+
+            HalfwordImm::strh_post(rt.reg_no, rn.base_reg_no, rn.disp as i16).into()
+        },
+        _ => panic!("Invalid operand combination to strh instruction.")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
 /// STUR - store a value in a register at a memory address
 pub fn stur(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
     let bytes: [u8; 4] = match (rt, rn) {
@@ -1099,6 +1189,21 @@ mod tests {
     }
 
     #[test]
+    fn test_ldrh() {
+        check_bytes("6a194079", |cb| ldrh(cb, W10, A64Opnd::new_mem(64, X11, 12)));
+    }
+
+    #[test]
+    fn test_ldrh_pre() {
+        check_bytes("6acd4078", |cb| ldrh_pre(cb, W10, A64Opnd::new_mem(64, X11, 12)));
+    }
+
+    #[test]
+    fn test_ldrh_post() {
+        check_bytes("6ac54078", |cb| ldrh_post(cb, W10, A64Opnd::new_mem(64, X11, 12)));
+    }
+
+    #[test]
     fn test_ldur_memory() {
         check_bytes("20b047f8", |cb| ldur(cb, X0, A64Opnd::new_mem(64, X1, 123)));
     }
@@ -1224,6 +1329,21 @@ mod tests {
     }
 
     #[test]
+    fn test_strh() {
+        check_bytes("6a190079", |cb| strh(cb, W10, A64Opnd::new_mem(64, X11, 12)));
+    }
+
+    #[test]
+    fn test_strh_pre() {
+        check_bytes("6acd0078", |cb| strh_pre(cb, W10, A64Opnd::new_mem(64, X11, 12)));
+    }
+
+    #[test]
+    fn test_strh_post() {
+        check_bytes("6ac50078", |cb| strh_post(cb, W10, A64Opnd::new_mem(64, X11, 12)));
+    }
+
+    #[test]
     fn test_stur() {
         check_bytes("6a0108f8", |cb| stur(cb, X10, A64Opnd::new_mem(64, X11, 128)));
     }