YJIT: fix 32 and 16 bit register store (#6840)

* Fix 32 and 16 bit register store in YJIT

Co-Authored-By: Takashi Kokubun <takashikkbn@gmail.com>

* Remove an unnecessary diff

* Reuse an rm_num_bits result

* Use u16::MAX instead

* Update the link

Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

* Just use sturh for 16 bits

Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>
Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

2 files changed, 22 insertions, 1 deletions

diff --git a/yjit/src/asm/arm64/inst/load_store.rs b/yjit/src/asm/arm64/inst/load_store.rs
index d08398380a..b5c8a3c294 100644
--- a/yjit/src/asm/arm64/inst/load_store.rs
+++ b/yjit/src/asm/arm64/inst/load_store.rs
@@ -118,6 +118,12 @@ impl LoadStore {
     pub fn stur(rt: u8, rn: u8, imm9: i16, num_bits: u8) -> Self {
         Self { rt, rn, idx: Index::None, imm9, opc: Opc::STR, size: num_bits.into() }
     }
+
+    /// STURH (store register, halfword, unscaled)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/STURH--Store-Register-Halfword--unscaled--?lang=en
+    pub fn sturh(rt: u8, rn: u8, imm9: i16) -> Self {
+        Self { rt, rn, idx: Index::None, imm9, opc: Opc::STR, size: Size::Size16 }
+    }
 }
 
 /// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/Loads-and-Stores?lang=en
diff --git a/yjit/src/asm/arm64/mod.rs b/yjit/src/asm/arm64/mod.rs
index d463613943..04660e580e 100644
--- a/yjit/src/asm/arm64/mod.rs
+++ b/yjit/src/asm/arm64/mod.rs
@@ -903,7 +903,7 @@ pub fn strh_post(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
 pub fn stur(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
     let bytes: [u8; 4] = match (rt, rn) {
         (A64Opnd::Reg(rt), A64Opnd::Mem(rn)) => {
-            assert!(rt.num_bits == rn.num_bits, "Expected registers to be the same size");
+            assert!(rn.num_bits == 32 || rn.num_bits == 64);
             assert!(mem_disp_fits_bits(rn.disp), "Expected displacement to be 9 bits or less");
 
             LoadStore::stur(rt.reg_no, rn.base_reg_no, rn.disp as i16, rt.num_bits).into()
@@ -914,6 +914,21 @@ pub fn stur(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
     cb.write_bytes(&bytes);
 }
 
+/// STURH - store a value in a register at a memory address
+pub fn sturh(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
+    let bytes: [u8; 4] = match (rt, rn) {
+        (A64Opnd::Reg(rt), A64Opnd::Mem(rn)) => {
+            assert!(rn.num_bits == 16);
+            assert!(mem_disp_fits_bits(rn.disp), "Expected displacement to be 9 bits or less");
+
+            LoadStore::sturh(rt.reg_no, rn.base_reg_no, rn.disp as i16).into()
+        },
+        _ => panic!("Invalid operand combination to stur instruction.")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
 /// SUB - subtract rm from rn, put the result in rd, don't update flags
 pub fn sub(cb: &mut CodeBlock, rd: A64Opnd, rn: A64Opnd, rm: A64Opnd) {
     let bytes: [u8; 4] = match (rd, rn, rm) {