ZJIT: Allow Store with 8-bit Opnd::Mem

3 files changed, 38 insertions, 6 deletions

diff --git a/zjit/src/asm/arm64/inst/load_store.rs b/zjit/src/asm/arm64/inst/load_store.rs
index e27909ae35..e78edd5675 100644
--- a/zjit/src/asm/arm64/inst/load_store.rs
+++ b/zjit/src/asm/arm64/inst/load_store.rs
@@ -124,6 +124,12 @@ impl LoadStore {
     pub fn sturh(rt: u8, rn: u8, imm9: i16) -> Self {
         Self { rt, rn, idx: Index::None, imm9, opc: Opc::STR, size: Size::Size16 }
     }
+
+    /// STURB (store register, byte, unscaled)
+    /// <https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/STURH--Store-Register-Halfword--unscaled--?lang=en>
+    pub fn sturb(rt: u8, rn: u8, imm9: i16) -> Self {
+        Self { rt, rn, idx: Index::None, imm9, opc: Opc::STR, size: Size::Size8 }
+    }
 }
 
 /// <https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/Loads-and-Stores?lang=en>
diff --git a/zjit/src/asm/arm64/mod.rs b/zjit/src/asm/arm64/mod.rs
index ffb83cf16a..a445911731 100644
--- a/zjit/src/asm/arm64/mod.rs
+++ b/zjit/src/asm/arm64/mod.rs
@@ -1026,7 +1026,21 @@ pub fn sturh(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
 
             LoadStore::sturh(rt.reg_no, rn.base_reg_no, rn.disp as i16).into()
         },
-        _ => panic!("Invalid operand combination to stur instruction.")
+        _ => panic!("Invalid operand combination to sturh instruction: {rt:?}, {rn:?}")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
+pub fn sturb(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
+    let bytes: [u8; 4] = match (rt, rn) {
+        (A64Opnd::Reg(rt), A64Opnd::Mem(rn)) => {
+            assert!(rn.num_bits == 8);
+            assert!(mem_disp_fits_bits(rn.disp), "Expected displacement {} to be 9 bits or less", rn.disp);
+
+            LoadStore::sturb(rt.reg_no, rn.base_reg_no, rn.disp as i16).into()
+        },
+        _ => panic!("Invalid operand combination to sturb instruction: {rt:?}, {rn:?}")
     };
 
     cb.write_bytes(&bytes);
diff --git a/zjit/src/backend/arm64/mod.rs b/zjit/src/backend/arm64/mod.rs
index b3827ae75d..428d4bff77 100644
--- a/zjit/src/backend/arm64/mod.rs
+++ b/zjit/src/backend/arm64/mod.rs
@@ -1298,7 +1298,8 @@ impl Assembler {
                     match dest.rm_num_bits() {
                         64 | 32 => stur(cb, src, dest.into()),
                         16 => sturh(cb, src, dest.into()),
-                        num_bits => panic!("unexpected dest num_bits: {} (src: {:#?}, dest: {:#?})", num_bits, src, dest),
+                        8 => sturb(cb, src, dest.into()),
+                        num_bits => panic!("unexpected dest num_bits: {} (src: {:?}, dest: {:?})", num_bits, src, dest),
                     }
                 },
                 Insn::Load { opnd, out } |
@@ -1675,6 +1676,7 @@ mod tests {
     static TEMP_REGS: [Reg; 5] = [X1_REG, X9_REG, X10_REG, X14_REG, X15_REG];
 
     fn setup_asm() -> (Assembler, CodeBlock) {
+        crate::options::rb_zjit_prepare_options(); // Allow `get_option!` in Assembler
         (Assembler::new(), CodeBlock::new_dummy())
     }
 
@@ -2601,8 +2603,21 @@ mod tests {
     }
 
     #[test]
+    fn test_store_spilled_byte() {
+        let (mut asm, mut cb) = setup_asm();
+
+        asm.store(Opnd::mem(8, C_RET_OPND, 0), Opnd::mem(8, C_RET_OPND, 8));
+        asm.compile_with_num_regs(&mut cb, 0); // spill every VReg
+
+        assert_disasm_snapshot!(cb.disasm(), @r"
+        0x0: ldurb w16, [x0, #8]
+        0x4: sturb w16, [x0]
+        ");
+        assert_snapshot!(cb.hexdump(), @"1080403810000038");
+    }
+
+    #[test]
     fn test_ccall_resolve_parallel_moves_no_cycle() {
-        crate::options::rb_zjit_prepare_options();
         let (mut asm, mut cb) = setup_asm();
 
         asm.ccall(0 as _, vec![
@@ -2620,7 +2635,6 @@ mod tests {
 
     #[test]
     fn test_ccall_resolve_parallel_moves_single_cycle() {
-        crate::options::rb_zjit_prepare_options();
         let (mut asm, mut cb) = setup_asm();
 
         // x0 and x1 form a cycle
@@ -2643,7 +2657,6 @@ mod tests {
 
     #[test]
     fn test_ccall_resolve_parallel_moves_two_cycles() {
-        crate::options::rb_zjit_prepare_options();
         let (mut asm, mut cb) = setup_asm();
 
         // x0 and x1 form a cycle, and x2 and rcx form another cycle
@@ -2670,7 +2683,6 @@ mod tests {
 
     #[test]
     fn test_ccall_resolve_parallel_moves_large_cycle() {
-        crate::options::rb_zjit_prepare_options();
         let (mut asm, mut cb) = setup_asm();
 
         // x0, x1, and x2 form a cycle