YJIT: Properly preserve register mapping in cpush_all() and cpop_all()

Previously, cpop_all() did not in fact restore the register mapping
state since it was effectively doing a no-op
`self.ctx.set_reg_mapping(self.ctx.get_reg_mapping())`. This desync in
bookkeeping led to issues with the --yjit-dump-insns option because
print_str() used to use cpush_all() and cpop_all().

Author: XrXr

yjit/src/backend/arm64/mod.rs

@@ -1419,7 +1419,7 @@ mod tests {
     fn test_emit_cpop_all() {
         let (mut asm, mut cb) = setup_asm();
 
-        asm.cpop_all();
+        asm.cpop_all(crate::core::RegMapping::default());
         asm.compile_with_num_regs(&mut cb, 0);
     }
 

yjit/src/backend/ir.rs

@@ -1824,12 +1824,12 @@ impl Assembler {
         out
     }
 
-    pub fn cpop_all(&mut self) {
+    pub fn cpop_all(&mut self, reg_mapping: RegMapping) {
         self.push_insn(Insn::CPopAll);
 
         // Re-enable ccall's RegMappings assertion disabled by cpush_all.
         // cpush_all + cpop_all preserve all stack temp registers, so it's safe.
-        self.set_reg_mapping(self.ctx.get_reg_mapping());
+        self.set_reg_mapping(reg_mapping);
     }
 
     pub fn cpop_into(&mut self, opnd: Opnd) {
@@ -1840,14 +1840,16 @@ impl Assembler {
         self.push_insn(Insn::CPush(opnd));
     }
 
-    pub fn cpush_all(&mut self) {
+    pub fn cpush_all(&mut self) -> RegMapping {
         self.push_insn(Insn::CPushAll);
 
         // Mark all temps as not being in registers.
         // Temps will be marked back as being in registers by cpop_all.
         // We assume that cpush_all + cpop_all are used for C functions in utils.rs
         // that don't require spill_regs for GC.
+        let mapping = self.ctx.get_reg_mapping();
         self.set_reg_mapping(RegMapping::default());
+        mapping
     }
 
     pub fn cret(&mut self, opnd: Opnd) {

yjit/src/codegen.rs

@@ -8314,7 +8314,7 @@ fn gen_send_iseq(
     // We also do this after spill_regs() to avoid doubly spilling the same thing on asm.ccall().
     if get_option!(gen_stats) {
         // Protect caller-saved registers in case they're used for arguments
-        asm.cpush_all();
+        let mapping = asm.cpush_all();
 
         // Assemble the ISEQ name string
         let name_str = get_iseq_name(iseq);
@@ -8324,7 +8324,7 @@ fn gen_send_iseq(
 
         // Increment the counter for this cfunc
         asm.ccall(incr_iseq_counter as *const u8, vec![iseq_idx.into()]);
-        asm.cpop_all();
+        asm.cpop_all(mapping);
     }
 
     // The callee might change locals through Kernel#binding and other means.