YJIT: Spill a caller stack as late as possible (#7726)

2 files changed, 12 insertions, 11 deletions

diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index 710464e44a..2e85a660ec 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -1189,7 +1189,7 @@ impl Assembler
 
         // Every stack temp should have been spilled
         assert_eq!(self.get_reg_temps(), RegTemps::default());
-        self.ctx.set_reg_temps(self.get_reg_temps());
+        assert_eq!(self.ctx.get_reg_temps(), RegTemps::default());
     }
 
     /// Sets the out field on the various instructions that require allocated
@@ -1815,13 +1815,18 @@ impl Assembler {
         if self.get_reg_temps() != reg_temps {
             self.comment(&format!("reg_temps: {:08b} -> {:08b}", self.get_reg_temps().as_u8(), reg_temps.as_u8()));
             self.push_insn(Insn::RegTemps(reg_temps));
+            self.ctx.set_reg_temps(self.get_reg_temps());
         }
     }
 
     /// Spill a stack temp from a register to the stack
     pub fn spill_temp(&mut self, opnd: Opnd) {
-        assert!(self.get_reg_temps().get(opnd.stack_idx()));
-        self.push_insn(Insn::SpillTemp(opnd));
+        assert_eq!(self.get_reg_temps(), self.ctx.get_reg_temps());
+
+        if opnd.stack_idx() < MAX_REG_TEMPS && self.get_reg_temps().get(opnd.stack_idx()) {
+            self.push_insn(Insn::SpillTemp(opnd));
+            self.ctx.set_reg_temps(self.get_reg_temps());
+        }
     }
 
     pub fn store(&mut self, dest: Opnd, src: Opnd) {
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 03c4e313fd..7fc110e1dc 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -4841,6 +4841,9 @@ fn gen_push_frame(
         asm.store(Opnd::mem(64, sp, offset), rest_arg);
     }
 
+    // Spill stack temps to let the callee use them (must be done before changing SP)
+    asm.spill_temps();
+
     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
@@ -5079,7 +5082,6 @@ fn gen_send_cfunc(
         assert_ne!(0, unsafe { rb_IMEMO_TYPE_P(imemo_ci, imemo_callinfo) },
             "we assume all callinfos with kwargs are on the GC heap");
         let sp = asm.lea(asm.ctx.sp_opnd(0));
-        asm.spill_temps(); // for ccall
         let kwargs = asm.ccall(build_kwhash as *const u8, vec![imemo_ci.into(), sp]);
 
         // Replace the stack location at the start of kwargs with the new hash
@@ -5090,9 +5092,6 @@ fn gen_send_cfunc(
     // Copy SP because REG_SP will get overwritten
     let sp = asm.lea(asm.ctx.sp_opnd(0));
 
-    // Arguments must be spilled before popped from ctx
-    asm.spill_temps();
-
     // Pop the C function arguments from the stack (in the caller)
     asm.stack_pop((argc + 1).try_into().unwrap());
 
@@ -6061,7 +6060,7 @@ fn gen_send_iseq(
         // pushed onto the stack that represents the parameters that weren't
         // explicitly given a value and have a non-constant default.
         let unspec_opnd = VALUE::fixnum_from_usize(unspecified_bits).as_u64();
-        asm.spill_temps(); // avoid using a register for unspecified_bits
+        asm.spill_temp(asm.stack_opnd(-1)); // avoid using a register for unspecified_bits
         asm.mov(asm.stack_opnd(-1), unspec_opnd.into());
     }
 
@@ -6094,9 +6093,6 @@ fn gen_send_iseq(
         argc = lead_num;
     }
 
-    // Spill stack temps to let the callee use them
-    asm.spill_temps();
-
     // If we have a rest param and optional parameters,
     // we don't actually pass the rest parameter as an argument,
     // instead we set its value in the callee's locals