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

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]
     }

yjit/src/backend/ir.rs

@@ -570,19 +570,27 @@ 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();
+            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 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();
+            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);