ZJIT: Use Mem.num_bits in Mem split (#15177)

Fix the

```
write(2, "ruby: ZJIT has panicked. More info to follow...\n", 48) = 48
write(2, "\nthread '<unnamed>' panicked at zjit/src/backend/lir.rs:160:17:\nassertion failed: num_bits <= out_num_bits\n", 107) = 107
```

based on

```
#25 0x0000aaaaaae8fb14 in zjit::backend::lir::Opnd::mem (num_bits=64, base=..., disp=0) at zjit/src/backend/lir.rs:160
#26 zjit::backend::arm64::{impl#3}::arm64_split::split_memory_address (asm=<optimized out>,
    opnd=<error reading variable: Cannot access memory at address 0x0>) at zjit/src/backend/arm64/mod.rs:260
#27 zjit::backend::arm64::{impl#3}::arm64_split::split_load_operand (asm=<optimized out>, opnd=...) at zjit/src/backend/arm64/mod.rs:273
```

1 files changed, 44 insertions, 2 deletions

diff --git a/zjit/src/backend/arm64/mod.rs b/zjit/src/backend/arm64/mod.rs
index 79c07d42ad..013fd31583 100644
--- a/zjit/src/backend/arm64/mod.rs
+++ b/zjit/src/backend/arm64/mod.rs
@@ -256,8 +256,8 @@ impl Assembler {
                     if mem_disp_fits_bits(mem.disp) {
                         opnd
                     } else {
-                        let base = asm.lea(opnd);
-                        Opnd::mem(64, base, 0)
+                        let base = asm.lea(Opnd::Mem(Mem { num_bits: 64, ..mem }));
+                        Opnd::mem(mem.num_bits, base, 0)
                     }
                 },
                 _ => unreachable!("Can only split memory addresses.")
@@ -2735,4 +2735,46 @@ mod tests {
         ");
         assert_snapshot!(cb.hexdump(), @"300080d2b0831ff8af835ff8eff97fd3af831ff8a0835ff8");
     }
+
+    #[test]
+    fn test_split_load16_mem_mem_with_large_displacement() {
+        let (mut asm, mut cb) = setup_asm();
+
+        let _ = asm.load(Opnd::mem(16, C_RET_OPND, 0x200));
+        asm.compile(&mut cb).unwrap();
+
+        assert_disasm_snapshot!(cb.disasm(), @r"
+        0x0: add x0, x0, #0x200
+        0x4: ldurh w0, [x0]
+        ");
+        assert_snapshot!(cb.hexdump(), @"0000089100004078");
+    }
+
+    #[test]
+    fn test_split_load32_mem_mem_with_large_displacement() {
+        let (mut asm, mut cb) = setup_asm();
+
+        let _ = asm.load(Opnd::mem(32, C_RET_OPND, 0x200));
+        asm.compile(&mut cb).unwrap();
+
+        assert_disasm_snapshot!(cb.disasm(), @r"
+        0x0: add x0, x0, #0x200
+        0x4: ldur w0, [x0]
+        ");
+        assert_snapshot!(cb.hexdump(), @"00000891000040b8");
+    }
+
+    #[test]
+    fn test_split_load64_mem_mem_with_large_displacement() {
+        let (mut asm, mut cb) = setup_asm();
+
+        let _ = asm.load(Opnd::mem(64, C_RET_OPND, 0x200));
+        asm.compile(&mut cb).unwrap();
+
+        assert_disasm_snapshot!(cb.disasm(), @r"
+        0x0: add x0, x0, #0x200
+        0x4: ldur x0, [x0]
+        ");
+        assert_snapshot!(cb.hexdump(), @"00000891000040f8");
+    }
 }