Op::CPushAll and Op::CPopAll (https://github.com/Shopify/ruby/pull/317)

Instructions for pushing all caller-save registers and the flags so that
we can implement dump_insns.

4 files changed, 116 insertions, 21 deletions

diff --git a/yjit/src/asm/arm64/opnd.rs b/yjit/src/asm/arm64/opnd.rs
index 1738f0985c..e1f95979a9 100644
--- a/yjit/src/asm/arm64/opnd.rs
+++ b/yjit/src/asm/arm64/opnd.rs
@@ -88,6 +88,7 @@ impl A64Opnd {
     }
 }
 
+// argument registers
 pub const X0_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 0 };
 pub const X1_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 1 };
 pub const X2_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 2 };
@@ -95,15 +96,20 @@ pub const X3_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 3 };
 pub const X4_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 4 };
 pub const X5_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 5 };
 
+// caller-save registers
 pub const X9_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 9 };
 pub const X10_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 10 };
 pub const X11_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 11 };
 pub const X12_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 12 };
 pub const X13_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 13 };
+pub const X14_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 14 };
+pub const X15_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 15 };
 
-pub const X24_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 24 };
-pub const X25_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 25 };
-pub const X26_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 26 };
+// callee-save registers
+pub const X19_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 19 };
+pub const X20_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 20 };
+pub const X21_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 21 };
+pub const X22_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 22 };
 
 // 64-bit registers
 pub const X0: A64Opnd = A64Opnd::Reg(X0_REG);
@@ -120,19 +126,19 @@ pub const X10: A64Opnd = A64Opnd::Reg(X10_REG);
 pub const X11: A64Opnd = A64Opnd::Reg(X11_REG);
 pub const X12: A64Opnd = A64Opnd::Reg(X12_REG);
 pub const X13: A64Opnd = A64Opnd::Reg(X13_REG);
-pub const X14: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 14 });
-pub const X15: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 15 });
+pub const X14: A64Opnd = A64Opnd::Reg(X14_REG);
+pub const X15: A64Opnd = A64Opnd::Reg(X15_REG);
 pub const X16: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 16 });
 pub const X17: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 17 });
 pub const X18: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 18 });
-pub const X19: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 19 });
-pub const X20: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 20 });
-pub const X21: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 21 });
-pub const X22: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 22 });
+pub const X19: A64Opnd = A64Opnd::Reg(X19_REG);
+pub const X20: A64Opnd = A64Opnd::Reg(X20_REG);
+pub const X21: A64Opnd = A64Opnd::Reg(X21_REG);
+pub const X22: A64Opnd = A64Opnd::Reg(X22_REG);
 pub const X23: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 23 });
-pub const X24: A64Opnd = A64Opnd::Reg(X24_REG);
-pub const X25: A64Opnd = A64Opnd::Reg(X25_REG);
-pub const X26: A64Opnd = A64Opnd::Reg(X26_REG);
+pub const X24: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 24 });
+pub const X25: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 25 });
+pub const X26: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 26 });
 pub const X27: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 27 });
 pub const X28: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 28 });
 pub const X29: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 29 });
diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index f6429dbcea..a208eb6316 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -13,9 +13,9 @@ use crate::virtualmem::CodePtr;
 pub type Reg = A64Reg;
 
 // Callee-saved registers
-pub const _CFP: Opnd = Opnd::Reg(X24_REG);
-pub const _EC: Opnd = Opnd::Reg(X25_REG);
-pub const _SP: Opnd = Opnd::Reg(X26_REG);
+pub const _CFP: Opnd = Opnd::Reg(X19_REG);
+pub const _EC: Opnd = Opnd::Reg(X20_REG);
+pub const _SP: Opnd = Opnd::Reg(X21_REG);
 
 // C argument registers on this platform
 pub const _C_ARG_OPNDS: [Opnd; 6] = [
@@ -59,11 +59,15 @@ impl From<Opnd> for A64Opnd {
 impl Assembler
 {
     /// Get the list of registers from which we can allocate on this platform
-    pub fn get_alloc_regs() -> Vec<Reg>
-    {
+    pub fn get_alloc_regs() -> Vec<Reg> {
         vec![C_RET_REG, X12_REG]
     }
 
+    /// Get a list of all of the caller-save registers
+    pub fn get_caller_save_regs() -> Vec<Reg> {
+        vec![X9_REG, X10_REG, X11_REG, X12_REG, X13_REG, X14_REG, X15_REG]
+    }
+
     /// Split platform-specific instructions
     /// The transformations done here are meant to make our lives simpler in later
     /// stages of the compilation pipeline.
@@ -340,11 +344,28 @@ impl Assembler
             };
         }
 
+        /// Emit a push instruction for the given operand by adding to the stack
+        /// pointer and then storing the given value.
+        fn emit_push(cb: &mut CodeBlock, opnd: A64Opnd) {
+            add(cb, C_SP_REG, C_SP_REG, C_SP_STEP);
+            stur(cb, opnd, A64Opnd::new_mem(64, C_SP_REG, 0));
+        }
+
+        /// Emit a pop instruction into the given operand by loading the value
+        /// and then subtracting from the stack pointer.
+        fn emit_pop(cb: &mut CodeBlock, opnd: A64Opnd) {
+            ldur(cb, opnd, A64Opnd::new_mem(64, C_SP_REG, 0));
+            sub(cb, C_SP_REG, C_SP_REG, C_SP_STEP);
+        }
+
         // dbg!(&self.insns);
 
         // List of GC offsets
         let mut gc_offsets: Vec<u32> = Vec::new();
 
+        // A special scratch register for loading/storing system registers.
+        let mut sys_scratch = A64Opnd::Reg(X22_REG);
+
         // For each instruction
         for insn in &self.insns {
             match insn.op {
@@ -429,12 +450,30 @@ impl Assembler
                     };
                 },
                 Op::CPush => {
-                    add(cb, C_SP_REG, C_SP_REG, C_SP_STEP);
-                    stur(cb, insn.opnds[0].into(), A64Opnd::new_mem(64, C_SP_REG, 0));
+                    emit_push(cb, insn.opnds[0].into());
+                },
+                Op::CPushAll => {
+                    let regs = Assembler::get_caller_save_regs();
+
+                    for reg in regs {
+                        emit_push(cb, A64Opnd::Reg(reg));
+                    }
+
+                    mrs(cb, sys_scratch, SystemRegister::NZCV);
+                    emit_push(cb, sys_scratch);
                 },
                 Op::CPop => {
-                    ldur(cb, insn.opnds[0].into(), A64Opnd::new_mem(64, C_SP_REG, 0));
-                    sub(cb, C_SP_REG, C_SP_REG, C_SP_STEP);
+                    emit_pop(cb, insn.opnds[0].into());
+                },
+                Op::CPopAll => {
+                    let regs = Assembler::get_caller_save_regs();
+
+                    msr(cb, SystemRegister::NZCV, sys_scratch);
+                    emit_pop(cb, sys_scratch);
+
+                    for reg in regs.into_iter().rev() {
+                        emit_pop(cb, A64Opnd::Reg(reg));
+                    }
                 },
                 Op::CCall => {
                     let src_addr = cb.get_write_ptr().into_i64() + 4;
@@ -570,4 +609,20 @@ mod tests {
         let insns = cb.get_ptr(0).raw_ptr() as *const u32;
         assert_eq!(0x8b010003, unsafe { *insns });
     }
+
+    #[test]
+    fn test_emit_cpush_all() {
+        let (mut asm, mut cb) = setup_asm();
+
+        asm.cpush_all();
+        asm.compile_with_num_regs(&mut cb, 0);
+    }
+
+    #[test]
+    fn test_emit_cpop_all() {
+        let (mut asm, mut cb) = setup_asm();
+
+        asm.cpop_all();
+        asm.compile_with_num_regs(&mut cb, 0);
+    }
 }
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index 5758d72d43..dbc6464a9c 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -97,6 +97,11 @@ pub enum Op
     CPush,
     CPop,
 
+    // Push and pop all of the caller-save registers and the flags to/from the C
+    // stack
+    CPushAll,
+    CPopAll,
+
     // C function call with N arguments (variadic)
     CCall,
 
@@ -804,6 +809,8 @@ def_push_2_opnd!(and, Op::And);
 def_push_1_opnd!(not, Op::Not);
 def_push_1_opnd_no_out!(cpush, Op::CPush);
 def_push_1_opnd_no_out!(cpop, Op::CPop);
+def_push_0_opnd_no_out!(cpush_all, Op::CPushAll);
+def_push_0_opnd_no_out!(cpop_all, Op::CPopAll);
 def_push_1_opnd_no_out!(cret, Op::CRet);
 def_push_1_opnd!(load, Op::Load);
 def_push_1_opnd!(lea, Op::Lea);
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index 4fd30e7144..4e0a9dcf02 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -76,6 +76,14 @@ impl Assembler
         ]
     }
 
+    /// Get a list of all of the caller-save registers
+    pub fn get_caller_save_regs() -> Vec<Reg> {
+        vec![RAX_REG, RCX_REG, RDX_REG, RSI_REG, RDI_REG, R8_REG, R9_REG, R10_REG, R11_REG]
+
+        // Technically these are also caller-save: R12_REG, R13_REG, R14_REG,
+        // and R15_REG, but we don't use them so we don't include them here.
+    }
+
     /// Split IR instructions for the x86 platform
     fn x86_split(mut self) -> Assembler
     {
@@ -239,6 +247,25 @@ impl Assembler
                 Op::CPush => push(cb, insn.opnds[0].into()),
                 Op::CPop => pop(cb, insn.opnds[0].into()),
 
+                // Push and pop to the C stack all caller-save registers and the
+                // flags
+                Op::CPushAll => {
+                    let regs = Assembler::get_caller_save_regs();
+
+                    for reg in regs {
+                        push(cb, X86Opnd::Reg(reg));
+                    }
+                    pushfq(cb);
+                },
+                Op::CPopAll => {
+                    let regs = Assembler::get_caller_save_regs();
+
+                    popfq(cb);
+                    for reg in regs.into_iter().rev() {
+                        pop(cb, X86Opnd::Reg(reg));
+                    }
+                },
+
                 // C function call
                 Op::CCall => {
                     // Temporary