Add LiveReg IR instruction to fix stats leave exit code (https://github.com/Shopify/ruby/pull/341)

It allows for reserving a specific register and prevents the register
allocator from clobbering it. Without this
`./miniruby --yjit-stats --yjit-callthreshold=1 -e0` was crashing because
the counter incrementing code was clobbering RAX incorrectly.

4 files changed, 19 insertions, 7 deletions

diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index 9726a0f8f2..c6cd1b882c 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -792,6 +792,7 @@ impl Assembler
                 Op::CSelGE => {
                     csel(cb, insn.out.into(), insn.opnds[0].into(), insn.opnds[1].into(), Condition::GE);
                 }
+                Op::LiveReg => (), // just a reg alloc signal, no code
             };
         }
 
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index c55a8f609b..dc0e450df4 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -145,7 +145,10 @@ pub enum Op
     FrameSetup,
 
     /// Tear down the frame stack as necessary per the architecture.
-    FrameTeardown
+    FrameTeardown,
+
+    /// Take a specific register. Signal the register allocator to not use it.
+    LiveReg,
 }
 
 // Memory operand base
@@ -633,7 +636,6 @@ impl Assembler
             if let Some(reg_index) = reg_index {
                 assert_eq!(*pool & (1 << reg_index), 0);
                 *pool |= 1 << reg_index;
-                //return regs[reg_index];
             }
 
             return *reg;
@@ -713,7 +715,13 @@ impl Assembler
 
                 // Allocate a new register for this instruction
                 if out_reg == Opnd::None {
-                    out_reg = Opnd::Reg(alloc_reg(&mut pool, &regs))
+                    out_reg = if op == Op::LiveReg {
+                        // Allocate a specific register
+                        let reg = opnds[0].unwrap_reg();
+                        Opnd::Reg(take_reg(&mut pool, &regs, &reg))
+                    } else {
+                        Opnd::Reg(alloc_reg(&mut pool, &regs))
+                    }
                 }
             }
 
@@ -902,6 +910,7 @@ def_push_1_opnd_no_out!(cret, Op::CRet);
 def_push_1_opnd!(load, Op::Load);
 def_push_1_opnd!(load_sext, Op::LoadSExt);
 def_push_1_opnd!(lea, Op::Lea);
+def_push_1_opnd!(live_reg_opnd, Op::LiveReg);
 def_push_2_opnd_no_out!(store, Op::Store);
 def_push_2_opnd_no_out!(mov, Op::Mov);
 def_push_2_opnd_no_out!(cmp, Op::Cmp);
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index 7b84e62134..5bae5c7f29 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -448,7 +448,8 @@ impl Assembler
                 Op::CSelGE => {
                     mov(cb, insn.out.into(), insn.opnds[0].into());
                     cmovl(cb, insn.out.into(), insn.opnds[1].into());
-                },
+                }
+                Op::LiveReg => (), // just a reg alloc signal, no code
 
                 // We want to keep the panic here because some instructions that
                 // we feed to the backend could get lowered into other
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index fa0394eed5..e122d67910 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -213,7 +213,7 @@ macro_rules! gen_counter_incr {
             let counter_opnd = Opnd::mem(64, ptr_reg, 0);
 
             // Increment and store the updated value
-            $asm.incr_counter(counter_opnd, 1.into() );
+            $asm.incr_counter(counter_opnd, 1.into());
         }
     };
 }
@@ -552,8 +552,9 @@ fn gen_leave_exit(ocb: &mut OutlinedCb) -> CodePtr {
     let code_ptr = ocb.get_write_ptr();
     let mut asm = Assembler::new();
 
-    // NOTE: gen_leave() fully reconstructs interpreter state and leaves the
+    // gen_leave() fully reconstructs interpreter state and leaves the
     // return value in C_RET_OPND before coming here.
+    let ret_opnd = asm.live_reg_opnd(C_RET_OPND);
 
     // Every exit to the interpreter should be counted
     gen_counter_incr!(asm, leave_interp_return);
@@ -564,7 +565,7 @@ fn gen_leave_exit(ocb: &mut OutlinedCb) -> CodePtr {
 
     asm.frame_teardown();
 
-    asm.cret(C_RET_OPND);
+    asm.cret(ret_opnd);
 
     asm.compile(ocb);