ZJIT: A64: Fix splitting for large memory displacements

On the ruby side, this fixes a crash for methods with 39 or more
parameters. We used to miscomp those entry points due to Insn::Lea
picking ADDS which cannot reference SP:

    # set method params: 40
    mov x0, #0xfee8
    movk x0, #0xffff, lsl #16
    movk x0, #0xffff, lsl #32
    movk x0, #0xffff, lsl #48
    adds x0, xzr, x0

Have Lea work for all i32 displacements and avoid involving the split
pass. Previously, direct use of Insn::Lea directly from the user (as
opposed to generated by the split pass for some memory operations)
wasn't split, so being able to handle the whole range in arm64_emit()
was implicitly required. Also, not going through split reduces register
pressure.

test/ruby/test_zjit.rb

@@ -826,6 +826,17 @@ class TestZJIT < Test::Unit::TestCase
     }
   end
 
+  def test_forty_param_method
+    # This used to a trigger a miscomp on A64 due
+    # to a memory displacement larger than 9 bits.
+    assert_compiles '1', %Q{
+      def foo(#{'_,' * 39} n40) = n40
+
+      foo(0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1)
+    }
+  end
+
+
   def test_opt_aref_with
     assert_compiles ':ok', %q{
       def aref_with(hash) = hash["key"]

zjit/src/backend/arm64/mod.rs

@@ -219,29 +219,6 @@ impl Assembler
     /// have no memory operands.
     fn arm64_split(mut self) -> Assembler
     {
-        /// When we're attempting to load a memory address into a register, the
-        /// displacement must fit into the maximum number of bits for an Op::Add
-        /// immediate. If it doesn't, we have to load the displacement into a
-        /// register first.
-        fn split_lea_operand(asm: &mut Assembler, opnd: Opnd) -> Opnd {
-            match opnd {
-                Opnd::Mem(Mem { base, disp, num_bits }) => {
-                    if disp >= 0 && ShiftedImmediate::try_from(disp as u64).is_ok() {
-                        asm.lea(opnd)
-                    } else {
-                        let disp = asm.load(Opnd::Imm(disp.into()));
-                        let reg = match base {
-                            MemBase::Reg(reg_no) => Opnd::Reg(Reg { reg_no, num_bits }),
-                            MemBase::VReg(idx) => Opnd::VReg { idx, num_bits }
-                        };
-
-                        asm.add(reg, disp)
-                    }
-                },
-                _ => unreachable!("Op::Lea only accepts Opnd::Mem operands.")
-            }
-        }
-
         /// When you're storing a register into a memory location or loading a
         /// memory location into a register, the displacement from the base
         /// register of the memory location must fit into 9 bits. If it doesn't,
@@ -252,7 +229,7 @@ impl Assembler
                     if mem_disp_fits_bits(mem.disp) {
                         opnd
                     } else {
-                        let base = split_lea_operand(asm, opnd);
+                        let base = asm.lea(opnd);
                         Opnd::mem(64, base, 0)
                     }
                 },
@@ -575,7 +552,7 @@ impl Assembler
                 },
                 Insn::IncrCounter { mem, value } => {
                     let counter_addr = match mem {
-                        Opnd::Mem(_) => split_lea_operand(asm, *mem),
+                        Opnd::Mem(_) => asm.lea(*mem),
                         _ => *mem
                     };
 
@@ -1113,22 +1090,24 @@ impl Assembler
                     }
                 },
                 Insn::Lea { opnd, out } => {
-                    let opnd: A64Opnd = opnd.into();
+                    let &Opnd::Mem(Mem { num_bits: _, base: MemBase::Reg(base_reg_no), disp }) = opnd else {
-
+                        panic!("Unexpected Insn::Lea operand in arm64_emit: {opnd:?}");
-                    match opnd {
-                        A64Opnd::Mem(mem) => {
-                            add(
-                                cb,
-                                out.into(),
-                                A64Opnd::Reg(A64Reg { reg_no: mem.base_reg_no, num_bits: 64 }),
-                                A64Opnd::new_imm(mem.disp.into())
-                            );
-                        },
-                        _ => {
-                            panic!("Op::Lea only accepts Opnd::Mem operands.");
-                        }
                     };
-                },
+                    let out: A64Opnd = out.into();
+                    let base_reg = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: base_reg_no });
+                    assert_ne!(31, out.unwrap_reg().reg_no, "Insn::Lea sp, [sp, #imm] not always encodable. Use add/sub instead.");
+
+                    if ShiftedImmediate::try_from(disp.unsigned_abs() as u64).is_ok() {
+                        // Use ADD/SUB if the displacement fits
+                        add(cb, out, base_reg, A64Opnd::new_imm(disp.into()));
+                    } else {
+                        // Use add_extended() to interpret reg_no=31 as sp
+                        // since the base register is never the zero register.
+                        // Careful! Only the first two operands can refer to sp.
+                        emit_load_value(cb, out, disp as u64);
+                        add_extended(cb, out, base_reg, out);
+                    };
+                }
                 Insn::LeaJumpTarget { out, target, .. } => {
                     if let Target::Label(label_idx) = target {
                         // Set output to the raw address of the label
@@ -1607,6 +1586,45 @@ mod tests {
         asm.compile_with_num_regs(&mut cb, 0);
     }
 
+    #[test]
+    fn test_emit_lea() {
+        let (mut asm, mut cb) = setup_asm();
+
+        // Test values that exercise various types of immediates.
+        //  - 9 bit displacement for Load/Store
+        //  - 12 bit shifted immediate
+        //  - bit mask immediates
+        for displacement in [i32::MAX, 0x10008, 0x1800, 0x208, -0x208, -0x1800, -0x1008, i32::MIN] {
+            let mem = Opnd::mem(64, NATIVE_STACK_PTR, displacement);
+            asm.lea_into(Opnd::Reg(X0_REG), mem);
+        }
+
+        asm.compile_with_num_regs(&mut cb, 0);
+        assert_disasm!(cb, "e07b40b2e063208b000180d22000a0f2e063208b000083d2e063208be0230891e02308d100009dd2e0ffbff2e0ffdff2e0fffff2e063208b00ff9dd2e0ffbff2e0ffdff2e0fffff2e063208be08361b2e063208b", "
+            0x0: orr x0, xzr, #0x7fffffff
+            0x4: add x0, sp, x0
+            0x8: mov x0, #8
+            0xc: movk x0, #1, lsl #16
+            0x10: add x0, sp, x0
+            0x14: mov x0, #0x1800
+            0x18: add x0, sp, x0
+            0x1c: add x0, sp, #0x208
+            0x20: sub x0, sp, #0x208
+            0x24: mov x0, #0xe800
+            0x28: movk x0, #0xffff, lsl #16
+            0x2c: movk x0, #0xffff, lsl #32
+            0x30: movk x0, #0xffff, lsl #48
+            0x34: add x0, sp, x0
+            0x38: mov x0, #0xeff8
+            0x3c: movk x0, #0xffff, lsl #16
+            0x40: movk x0, #0xffff, lsl #32
+            0x44: movk x0, #0xffff, lsl #48
+            0x48: add x0, sp, x0
+            0x4c: orr x0, xzr, #0xffffffff80000000
+            0x50: add x0, sp, x0
+        ");
+    }
+
     /*
     #[test]
     fn test_emit_lea_label() {