YJIT: Avoid creating a vector in get_temp_regs() (#8446)

* YJIT: Avoid creating a vector in get_temp_regs()

Co-authored-by: Alan Wu <alansi.xingwu@shopify.com>

* Remove unused import

---------

Co-authored-by: Alan Wu <alansi.xingwu@shopify.com>
Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

yjit/src/backend/arm64/mod.rs

@@ -184,6 +184,10 @@ fn emit_load_value(cb: &mut CodeBlock, rd: A64Opnd, value: u64) -> usize {
     }
 }
 
+/// List of registers that can be used for stack temps.
+/// These are caller-saved registers.
+pub static TEMP_REGS: [Reg; 5] = [X1_REG, X9_REG, X10_REG, X14_REG, X15_REG];
+
 impl Assembler
 {
     // Special scratch registers for intermediate processing.
@@ -192,10 +196,6 @@ impl Assembler
     const SCRATCH0: A64Opnd = A64Opnd::Reg(Assembler::SCRATCH_REG);
     const SCRATCH1: A64Opnd = A64Opnd::Reg(X17_REG);
 
-    /// List of registers that can be used for stack temps.
-    /// These are caller-saved registers.
-    pub const TEMP_REGS: [Reg; 5] = [X1_REG, X9_REG, X10_REG, X14_REG, X15_REG];
-
     /// 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
@@ -1636,7 +1636,7 @@ mod tests {
     fn test_replace_mov_with_ldur() {
         let (mut asm, mut cb) = setup_asm();
 
-        asm.mov(Opnd::Reg(Assembler::TEMP_REGS[0]), Opnd::mem(64, CFP, 8));
+        asm.mov(Opnd::Reg(TEMP_REGS[0]), Opnd::mem(64, CFP, 8));
         asm.compile_with_num_regs(&mut cb, 1);
 
         assert_disasm!(cb, "618240f8", {"
@@ -1648,8 +1648,8 @@ mod tests {
     fn test_not_split_mov() {
         let (mut asm, mut cb) = setup_asm();
 
-        asm.mov(Opnd::Reg(Assembler::TEMP_REGS[0]), Opnd::UImm(0xffff));
-        asm.mov(Opnd::Reg(Assembler::TEMP_REGS[0]), Opnd::UImm(0x10000));
+        asm.mov(Opnd::Reg(TEMP_REGS[0]), Opnd::UImm(0xffff));
+        asm.mov(Opnd::Reg(TEMP_REGS[0]), Opnd::UImm(0x10000));
         asm.compile_with_num_regs(&mut cb, 1);
 
         assert_disasm!(cb, "e1ff9fd2e10370b2", {"
@@ -1663,7 +1663,7 @@ mod tests {
         let (mut asm, mut cb) = setup_asm();
 
         let out = asm.csel_l(Qtrue.into(), Qfalse.into());
-        asm.mov(Opnd::Reg(Assembler::TEMP_REGS[0]), out);
+        asm.mov(Opnd::Reg(TEMP_REGS[0]), out);
         asm.compile_with_num_regs(&mut cb, 2);
 
         assert_disasm!(cb, "8b0280d20c0080d261b18c9a", {"

yjit/src/backend/ir.rs

@@ -14,11 +14,7 @@ use crate::core::{Context, RegTemps, MAX_REG_TEMPS};
 use crate::options::*;
 use crate::stats::*;
 
-#[cfg(target_arch = "x86_64")]
-use crate::backend::x86_64::*;
-
-#[cfg(target_arch = "aarch64")]
-use crate::backend::arm64::*;
+use crate::backend::current::*;
 
 pub const EC: Opnd = _EC;
 pub const CFP: Opnd = _CFP;
@@ -1028,10 +1024,9 @@ impl Assembler
     }
 
     /// Get the list of registers that can be used for stack temps.
-    pub fn get_temp_regs() -> Vec<Reg> {
+    pub fn get_temp_regs() -> &'static [Reg] {
         let num_regs = get_option!(num_temp_regs);
-        let mut regs = Self::TEMP_REGS.to_vec();
-        regs.drain(0..num_regs).collect()
+        &TEMP_REGS[0..num_regs]
     }
 
     /// Set a context for generating side exits

yjit/src/backend/mod.rs

@@ -4,5 +4,11 @@ pub mod x86_64;
 #[cfg(target_arch = "aarch64")]
 pub mod arm64;
 
+#[cfg(target_arch = "x86_64")]
+pub use x86_64 as current;
+
+#[cfg(target_arch = "aarch64")]
+pub use arm64 as current;
+
 pub mod ir;
 mod tests;

yjit/src/backend/x86_64/mod.rs

@@ -84,6 +84,9 @@ impl From<&Opnd> for X86Opnd {
     }
 }
 
+/// List of registers that can be used for stack temps.
+pub static TEMP_REGS: [Reg; 5] = [RSI_REG, RDI_REG, R8_REG, R9_REG, R10_REG];
+
 impl Assembler
 {
     // A special scratch register for intermediate processing.
@@ -91,8 +94,6 @@ impl Assembler
     pub const SCRATCH_REG: Reg = R11_REG;
     const SCRATCH0: X86Opnd = X86Opnd::Reg(Assembler::SCRATCH_REG);
 
-    /// List of registers that can be used for stack temps.
-    pub const TEMP_REGS: [Reg; 5] = [RSI_REG, RDI_REG, R8_REG, R9_REG, R10_REG];
 
     /// Get the list of registers from which we can allocate on this platform
     pub fn get_alloc_regs() -> Vec<Reg>

yjit/src/core.rs

@@ -2694,7 +2694,7 @@ pub fn gen_branch_stub_hit_trampoline(ocb: &mut OutlinedCb) -> CodePtr {
 /// Return registers to be pushed and popped on branch_stub_hit.
 /// The return value may include an extra register for x86 alignment.
 fn caller_saved_temp_regs() -> Vec<Opnd> {
-    let mut regs = Assembler::get_temp_regs();
+    let mut regs = Assembler::get_temp_regs().to_vec();
     if regs.len() % 2 == 1 {
         regs.push(*regs.last().unwrap()); // x86 alignment
     }

yjit/src/options.rs

@@ -1,5 +1,5 @@
 use std::ffi::CStr;
-use crate::backend::ir::Assembler;
+use crate::backend::current::TEMP_REGS;
 
 // Command-line options
 #[derive(Clone, PartialEq, Eq, Debug)]
@@ -156,7 +156,7 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
 
         ("temp-regs", _) => match opt_val.parse() {
             Ok(n) => {
-                assert!(n <= Assembler::TEMP_REGS.len(), "--yjit-temp-regs must be <= {}", Assembler::TEMP_REGS.len());
+                assert!(n <= TEMP_REGS.len(), "--yjit-temp-regs must be <= {}", TEMP_REGS.len());
                 unsafe { OPTIONS.num_temp_regs = n }
             }
             Err(_) => {