Fix register preservation around ccall

We need to push pairs of registers so that code is more compact.  Also
don't try to preserve the return value of a ccall if the live range is
dead

Author: tenderlove

zjit/src/backend/arm64/mod.rs

@@ -2741,17 +2741,16 @@ mod tests {
         0x4: mov x1, #2
         0x8: mov x2, #3
         0xc: mov x3, #4
-        0x10: mov x4, x0
-        0x14: stp x2, x1, [sp, #-0x10]!
-        0x18: stp x4, x3, [sp, #-0x10]!
-        0x1c: mov x16, #0
-        0x20: blr x16
-        0x24: ldp x4, x3, [sp], #0x10
-        0x28: ldp x2, x1, [sp], #0x10
-        0x2c: adds x4, x4, x1
-        0x30: adds x2, x2, x3
+        0x10: stp x1, x0, [sp, #-0x10]!
+        0x14: stp x3, x2, [sp, #-0x10]!
+        0x18: mov x16, #0
+        0x1c: blr x16
+        0x20: ldp x3, x2, [sp], #0x10
+        0x24: ldp x1, x0, [sp], #0x10
+        0x28: adds x0, x0, x1
+        0x2c: adds x0, x2, x3
         ");
-        assert_snapshot!(cb.hexdump(), @"200080d2410080d2620080d2830080d2e40300aae207bfa9e40fbfa9100080d200023fd6e40fc1a8e207c1a8840001ab420003ab");
+        assert_snapshot!(cb.hexdump(), @"200080d2410080d2620080d2830080d2e103bfa9e30bbfa9100080d200023fd6e30bc1a8e103c1a8000001ab400003ab");
     }
 
     #[test]
@@ -2776,20 +2775,19 @@ mod tests {
         0x8: mov x2, #3
         0xc: mov x3, #4
         0x10: mov x4, #5
-        0x14: mov x5, x0
-        0x18: stp x2, x1, [sp, #-0x10]!
-        0x1c: stp x4, x3, [sp, #-0x10]!
-        0x20: str x5, [sp, #-0x10]!
-        0x24: mov x16, #0
-        0x28: blr x16
-        0x2c: ldr x5, [sp], #0x10
-        0x30: ldp x4, x3, [sp], #0x10
-        0x34: ldp x2, x1, [sp], #0x10
-        0x38: adds x5, x5, x1
-        0x3c: adds x0, x2, x3
-        0x40: adds x2, x2, x4
+        0x14: stp x1, x0, [sp, #-0x10]!
+        0x18: stp x3, x2, [sp, #-0x10]!
+        0x1c: str x4, [sp, #-0x10]!
+        0x20: mov x16, #0
+        0x24: blr x16
+        0x28: ldr x4, [sp], #0x10
+        0x2c: ldp x3, x2, [sp], #0x10
+        0x30: ldp x1, x0, [sp], #0x10
+        0x34: adds x0, x0, x1
+        0x38: adds x0, x2, x3
+        0x3c: adds x0, x2, x4
         ");
-        assert_snapshot!(cb.hexdump(), @"200080d2410080d2620080d2830080d2a40080d2e50300aae207bfa9e40fbfa9e50f1ff8100080d200023fd6e50741f8e40fc1a8e207c1a8a50001ab400003ab420004ab");
+        assert_snapshot!(cb.hexdump(), @"200080d2410080d2620080d2830080d2a40080d2e103bfa9e30bbfa9e40f1ff8100080d200023fd6e40741f8e30bc1a8e103c1a8000001ab400003ab400004ab");
     }
 
     #[test]

zjit/src/backend/lir.rs

@@ -2452,6 +2452,11 @@ impl Assembler
             for (insn, insn_id) in old_insns.into_iter().zip(old_ids.into_iter()) {
                 if let Insn::CCall { opnds, out, start_marker, end_marker, fptr } = insn {
                     let insn_number = insn_id.map(|id| id.0).unwrap_or(0);
+                    let call_result_live = out.is_vreg()
+                        && intervals[out.vreg_idx().0]
+                            .range
+                            .end
+                            .is_some_and(|end| end > insn_number);
 
                     // Find survivors: intervals that survive this Call instruction
                     // We need to preserve the "surviving" registers past the ccall,
@@ -2467,19 +2472,26 @@ impl Assembler
                         .map(|interval| interval.id)
                         .collect();
 
-                    // Push all survivors on the stack
+                    let survivor_regs: Vec<Opnd> = survivors.iter()
+                        .map(|&s| match assignments[s].unwrap() {
+                            Allocation::Reg(n) => Opnd::Reg(ALLOC_REGS[n]),
+                            Allocation::Fixed(reg) => Opnd::Reg(reg),
+                            _ => unreachable!(),
+                        })
+                        .collect();
+                    let survivor_push_groups: Vec<Vec<Opnd>> = survivor_regs
+                        .chunks(2)
+                        .map(|group| group.to_vec())
+                        .collect();
+
+                    // Push all survivors on the stack, pairing adjacent pushes when possible.
                     let needs_alignment = cfg!(target_arch = "x86_64") && survivors.len() % 2 == 1;
-                    for &s in &survivors {
-                        let reg_n = match assignments[s].unwrap() {
-                            Allocation::Reg(n) => n,
-                            Allocation::Fixed(reg) => {
-                                new_insns.push(Insn::CPush(Opnd::Reg(reg)));
-                                new_ids.push(None);
-                                continue;
-                            }
+                    for group in &survivor_push_groups {
+                        match group.as_slice() {
+                            [left, right] => new_insns.push(Insn::CPushPair(*left, *right)),
+                            [reg] => new_insns.push(Insn::CPush(*reg)),
                             _ => unreachable!(),
-                        };
-                        new_insns.push(Insn::CPush(Opnd::Reg(ALLOC_REGS[reg_n])));
+                        }
                         new_ids.push(None);
                     }
                     // Maintain 16-byte stack alignment for x86_64
@@ -2542,41 +2554,42 @@ impl Assembler
                     }
 
                     if survivors.is_empty() {
-                        // No survivors to restore — move result directly to output
-                        let out = Self::rewritten_opnd(out, assignments);
-                        new_insns.push(Insn::Mov { dest: out, src: C_RET_OPND });
-                        new_ids.push(None);
+                        if call_result_live {
+                            // No survivors to restore — move result directly to output.
+                            let out = Self::rewritten_opnd(out, assignments);
+                            new_insns.push(Insn::Mov { dest: out, src: C_RET_OPND });
+                            new_ids.push(None);
+                        }
                     } else {
-                        // Save CCall result to scratch immediately, before pops
-                        // can clobber either C_RET or the output register.
-                        new_insns.push(Insn::Mov { dest: Opnd::Reg(SCRATCH_REG), src: C_RET_OPND });
-                        new_ids.push(None);
+                        if call_result_live {
+                            // Save CCall result to scratch immediately, before pops
+                            // can clobber either C_RET or the output register.
+                            new_insns.push(Insn::Mov { dest: Opnd::Reg(SCRATCH_REG), src: C_RET_OPND });
+                            new_ids.push(None);
+                        }
 
                         // Pop alignment padding (if needed)
                         if needs_alignment {
                             new_insns.push(Insn::CPopInto(Opnd::Reg(ALLOC_REGS[0])));
                             new_ids.push(None);
                         }
 
-                        // Restore all survivors
-                        for &s in survivors.iter().rev() {
-                            let reg_n = match assignments[s].unwrap() {
-                                Allocation::Reg(n) => n,
-                                Allocation::Fixed(reg) => {
-                                    new_insns.push(Insn::CPopInto(Opnd::Reg(reg)));
-                                    new_ids.push(None);
-                                    continue;
-                                }
+                        // Restore all survivors in reverse stack order, pairing adjacent pops when possible.
+                        for group in survivor_push_groups.iter().rev() {
+                            match group.as_slice() {
+                                [left, right] => new_insns.push(Insn::CPopPairInto(*right, *left)),
+                                [reg] => new_insns.push(Insn::CPopInto(*reg)),
                                 _ => unreachable!(),
-                            };
-                            new_insns.push(Insn::CPopInto(Opnd::Reg(ALLOC_REGS[reg_n])));
+                            }
                             new_ids.push(None);
                         }
 
-                        // Move result from scratch to output AFTER all pops
-                        let out = Self::rewritten_opnd(out, assignments);
-                        new_insns.push(Insn::Mov { dest: out, src: Opnd::Reg(SCRATCH_REG) });
-                        new_ids.push(None);
+                        if call_result_live {
+                            // Move result from scratch to output AFTER all pops.
+                            let out = Self::rewritten_opnd(out, assignments);
+                            new_insns.push(Insn::Mov { dest: out, src: Opnd::Reg(SCRATCH_REG) });
+                            new_ids.push(None);
+                        }
                     }
                 } else {
                     new_insns.push(insn);