Exclude X0 (C_RET_REG) from allocatable registers on arm (https://github.com/Shopify/ruby/pull/319)

* Exclude X0 (C_RET_REG) from allocatable registers on arm

* Add another small test snippett

2 files changed, 20 insertions, 9 deletions

diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index a208eb6316..94da426ba3 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -58,12 +58,15 @@ impl From<Opnd> for A64Opnd {
 
 impl Assembler
 {
-    /// Get the list of registers from which we can allocate on this platform
+    /// Get the list of registers from which we will allocate on this platform
+    /// These are caller-saved registers
+    /// Note: we intentionally exclude C_RET_REG (X0) from this list
+    /// because of the way it's used in gen_leave() and gen_leave_exit()
     pub fn get_alloc_regs() -> Vec<Reg> {
-        vec![C_RET_REG, X12_REG]
+        vec![X11_REG, X12_REG]
     }
 
-    /// Get a list of all of the caller-save registers
+    /// Get a list of all of the caller-saved registers
     pub fn get_caller_save_regs() -> Vec<Reg> {
         vec![X9_REG, X10_REG, X11_REG, X12_REG, X13_REG, X14_REG, X15_REG]
     }
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index dbc6464a9c..89c12456a2 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -570,18 +570,26 @@ impl Assembler
 
         // Allocate a specific register
         fn take_reg(pool: &mut u32, regs: &Vec<Reg>, reg: &Reg) -> Reg {
-            let reg_index = regs.iter().position(|elem| elem.reg_no == reg.reg_no).unwrap();
-            assert_eq!(*pool & (1 << reg_index), 0);
-            *pool |= 1 << reg_index;
-            return regs[reg_index];
+            let reg_index = regs.iter().position(|elem| elem.reg_no == reg.reg_no);
+
+            if let Some(reg_index) = reg_index {
+                assert_eq!(*pool & (1 << reg_index), 0);
+                *pool |= 1 << reg_index;
+                //return regs[reg_index];
+            }
+
+            return *reg;
         }
 
         // Mutate the pool bitmap to indicate that the given register is being
         // returned as it is no longer used by the instruction that previously
         // held it.
         fn dealloc_reg(pool: &mut u32, regs: &Vec<Reg>, reg: &Reg) {
-            let reg_index = regs.iter().position(|elem| elem.reg_no == reg.reg_no).unwrap();
-            *pool &= !(1 << reg_index);
+            let reg_index = regs.iter().position(|elem| elem.reg_no == reg.reg_no);
+
+            if let Some(reg_index) = reg_index {
+                *pool &= !(1 << reg_index);
+            }
         }
 
         let live_ranges: Vec<usize> = std::mem::take(&mut self.live_ranges);