YJIT: Fix `yield` into block with >=30 locals on ARM

It's a register spill issue. Fix by moving the Qnil fill snippet to
after registers are released.

[Bug #19299]

1 files changed, 14 insertions, 15 deletions

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index e12a29240d..0347f6fc89 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -4455,19 +4455,6 @@ fn gen_push_frame(
     };
     asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -2), specval);
 
-    // Arm requires another register to load the immediate value of Qnil before storing it.
-    // So doing this after releasing the register for specval to avoid register spill.
-    let num_locals = frame.local_size;
-    if num_locals > 0 {
-        asm.comment("initialize locals");
-
-        // Initialize local variables to Qnil
-        for i in 0..num_locals {
-            let offs = (SIZEOF_VALUE as i32) * (i - num_locals - 3);
-            asm.store(Opnd::mem(64, sp, offs), Qnil.into());
-        }
-    }
-
     // Write env flags at sp[-1]
     // sp[-1] = frame_type;
     asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -1), frame.frame_type.into());
@@ -4505,6 +4492,20 @@ fn gen_push_frame(
     asm.mov(cfp_opnd(RUBY_OFFSET_CFP_SELF), frame.recv);
     asm.mov(cfp_opnd(RUBY_OFFSET_CFP_BLOCK_CODE), 0.into());
 
+    // This Qnil fill snippet potentially requires 2 more registers on Arm, one for Qnil and
+    // another for calculating the address in case there are a lot of local variables. So doing
+    // this after releasing the register for specval and the receiver to avoid register spill.
+    let num_locals = frame.local_size;
+    if num_locals > 0 {
+        asm.comment("initialize locals");
+
+        // Initialize local variables to Qnil
+        for i in 0..num_locals {
+            let offs = (SIZEOF_VALUE as i32) * (i - num_locals - 3);
+            asm.store(Opnd::mem(64, sp, offs), Qnil.into());
+        }
+    }
+
     if set_sp_cfp {
         // Saving SP before calculating ep avoids a dependency on a register
         // However this must be done after referencing frame.recv, which may be SP-relative
@@ -4778,8 +4779,6 @@ fn gen_return_branch(
     }
 }
 
-
-
 /// Pushes arguments from an array to the stack that are passed with a splat (i.e. *args)
 /// It optimistically compiles to a static size that is the exact number of arguments
 /// needed for the function.