Move backend tests to their own file

5 files changed, 238 insertions, 228 deletions

diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index 8685117c5f..be67e2384d 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -36,7 +36,10 @@ impl Assembler
     /// Get the list of registers from which we can allocate on this platform
     pub fn get_scratch_regs() -> Vec<Reg>
     {
-        vec![X12_REG, X13_REG]
+        vec![
+            X12_REG,
+            X13_REG
+        ]
     }
 
     /// Split platform-specific instructions
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index 1f6307db9e..4f1aafef99 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -604,8 +604,16 @@ impl Assembler
     ///       compiling multiple blocks at a time?
     pub fn compile(self, cb: &mut CodeBlock) -> Vec<u32>
     {
-        let scratch_regs = Self::get_scratch_regs();
-        self.compile_with_regs(cb, scratch_regs)
+        let alloc_regs = Self::get_alloc_regs();
+        self.compile_with_regs(cb, alloc_regs)
+    }
+
+    /// Compile with a limited number of registers
+    pub fn compile_with_num_regs(self, cb: &mut CodeBlock, num_regs: usize) -> Vec<u32>
+    {
+        let mut alloc_regs = Self::get_alloc_regs();
+        let alloc_regs = alloc_regs.drain(0..num_regs).collect();
+        self.compile_with_regs(cb, alloc_regs)
     }
 }
 
@@ -741,225 +749,3 @@ impl Context
         self.stack_push_mapping((mapping, temp_type)).into()
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::cruby::*;
-    use crate::core::*;
-    use InsnOpnd::*;
-
-    // Test that this function type checks
-    fn gen_dup(
-        ctx: &mut Context,
-        asm: &mut Assembler,
-    ) {
-        let dup_val = ctx.ir_stack_pop(0);
-        let (mapping, tmp_type) = ctx.get_opnd_mapping(StackOpnd(0));
-
-        let loc0 = ctx.ir_stack_push_mapping((mapping, tmp_type));
-        asm.mov(loc0, dup_val);
-    }
-
-    fn guard_object_is_heap(
-        asm: &mut Assembler,
-        object_opnd: Opnd,
-        ctx: &mut Context,
-        side_exit: CodePtr,
-    ) {
-        asm.comment("guard object is heap");
-
-        // Test that the object is not an immediate
-        asm.test(object_opnd.clone(), Opnd::UImm(RUBY_IMMEDIATE_MASK as u64));
-        asm.jnz(Target::CodePtr(side_exit));
-
-        // Test that the object is not false or nil
-        asm.cmp(object_opnd.clone(), Opnd::UImm(Qnil.into()));
-        asm.jbe(Target::CodePtr(side_exit));
-    }
-
-    #[test]
-    fn test_add() {
-        let mut asm = Assembler::new();
-        let out = asm.add(SP, Opnd::UImm(1));
-        asm.add(out, Opnd::UImm(2));
-    }
-
-    #[test]
-    fn test_split_loads() {
-        let mut asm = Assembler::new();
-
-        let regs = Assembler::get_scratch_regs();
-
-        asm.add(
-            Opnd::mem(64, Opnd::Reg(regs[0]), 0),
-            Opnd::mem(64, Opnd::Reg(regs[1]), 0)
-        );
-
-        let result = asm.split_loads();
-        assert_eq!(result.insns.len(), 2);
-        assert_eq!(result.insns[0].op, Op::Load);
-    }
-
-    #[test]
-    fn test_alloc_regs() {
-        let mut asm = Assembler::new();
-
-        // Get the first output that we're going to reuse later.
-        let out1 = asm.add(EC, Opnd::UImm(1));
-
-        // Pad some instructions in to make sure it can handle that.
-        asm.add(EC, Opnd::UImm(2));
-
-        // Get the second output we're going to reuse.
-        let out2 = asm.add(EC, Opnd::UImm(3));
-
-        // Pad another instruction.
-        asm.add(EC, Opnd::UImm(4));
-
-        // Reuse both the previously captured outputs.
-        asm.add(out1, out2);
-
-        // Now get a third output to make sure that the pool has registers to
-        // allocate now that the previous ones have been returned.
-        let out3 = asm.add(EC, Opnd::UImm(5));
-        asm.add(out3, Opnd::UImm(6));
-
-        // Here we're going to allocate the registers.
-        let result = asm.alloc_regs(Assembler::get_scratch_regs());
-
-        // Now we're going to verify that the out field has been appropriately
-        // updated for each of the instructions that needs it.
-        let regs = Assembler::get_scratch_regs();
-        assert_eq!(result.insns[0].out, Opnd::Reg(regs[0]));
-        assert_eq!(result.insns[2].out, Opnd::Reg(regs[1]));
-        assert_eq!(result.insns[5].out, Opnd::Reg(regs[0]));
-    }
-
-    fn setup_asm(num_regs: usize) -> (Assembler, CodeBlock, Vec<Reg>) {
-        let mut regs = Assembler::get_scratch_regs();
-
-        return (
-            Assembler::new(),
-            CodeBlock::new_dummy(1024),
-            regs.drain(0..num_regs).collect()
-        );
-    }
-
-    // Test full codegen pipeline
-    #[test]
-    fn test_compile()
-    {
-        let (mut asm, mut cb, regs) = setup_asm(1);
-
-        let out = asm.add(Opnd::Reg(regs[0]), Opnd::UImm(2));
-        asm.add(out, Opnd::UImm(2));
-
-        asm.compile(&mut cb);
-    }
-
-    // Test memory-to-memory move
-    #[test]
-    fn test_mov_mem2mem()
-    {
-        let (mut asm, mut cb, regs) = setup_asm(1);
-
-        asm.comment("check that comments work too");
-        asm.mov(Opnd::mem(64, SP, 0), Opnd::mem(64, SP, 8));
-
-        asm.compile_with_regs(&mut cb, regs);
-    }
-
-    // Test load of register into new register
-    #[test]
-    fn test_load_reg()
-    {
-        let (mut asm, mut cb, regs) = setup_asm(1);
-
-        let out = asm.load(SP);
-        asm.mov(Opnd::mem(64, SP, 0), out);
-
-        asm.compile_with_regs(&mut cb, regs);
-    }
-
-    // Multiple registers needed and register reuse
-    #[test]
-    fn test_reuse_reg()
-    {
-        let (mut asm, mut cb, regs) = setup_asm(2);
-
-        let v0 = asm.add(Opnd::mem(64, SP, 0), Opnd::UImm(1));
-        let v1 = asm.add(Opnd::mem(64, SP, 8), Opnd::UImm(1));
-        let v2 = asm.add(v0, Opnd::UImm(1));
-        asm.add(v0, v2);
-
-        asm.compile_with_regs(&mut cb, regs);
-    }
-
-    // 64-bit values can't be written directly to memory,
-    // need to be split into one or more register movs first
-    #[test]
-    fn test_store_u64()
-    {
-        let (mut asm, mut cb, regs) = setup_asm(1);
-        asm.store(Opnd::mem(64, SP, 0), u64::MAX.into());
-        asm.compile_with_regs(&mut cb, regs);
-    }
-
-    // Use instruction output as base register for memory operand
-    #[test]
-    fn test_base_insn_out()
-    {
-        let (mut asm, mut cb, regs) = setup_asm(1);
-
-        // Load the pointer into a register
-        let ptr_reg = asm.load(Opnd::const_ptr(0 as *const u8));
-        let counter_opnd = Opnd::mem(64, ptr_reg, 0);
-
-        // Increment and store the updated value
-        asm.incr_counter(counter_opnd, 1.into() );
-
-        asm.compile_with_regs(&mut cb, regs);
-    }
-
-    #[test]
-    fn test_c_call()
-    {
-        extern "sysv64" fn dummy_c_fun(v0: usize, v1: usize)
-        {
-        }
-
-        let (mut asm, mut cb, regs) = setup_asm(2);
-
-        asm.ccall(
-            dummy_c_fun as *const u8,
-            vec![Opnd::mem(64, SP, 0), Opnd::UImm(1)]
-        );
-
-        asm.compile_with_regs(&mut cb, regs);
-    }
-
-    #[test]
-    fn test_lea_ret()
-    {
-        let (mut asm, mut cb, regs) = setup_asm(1);
-
-        let addr = asm.lea(Opnd::mem(64, SP, 0));
-        asm.cret(addr);
-
-        asm.compile_with_regs(&mut cb, regs);
-    }
-
-    #[test]
-    fn test_jcc_label()
-    {
-        let (mut asm, mut cb, regs) = setup_asm(1);
-
-        let label = asm.new_label("foo");
-        asm.cmp(EC, EC);
-        asm.je(label);
-        asm.write_label(label);
-
-        asm.compile_with_regs(&mut cb, regs);
-    }
-}
diff --git a/yjit/src/backend/mod.rs b/yjit/src/backend/mod.rs
index a83fe4f69e..0841c9ffa5 100644
--- a/yjit/src/backend/mod.rs
+++ b/yjit/src/backend/mod.rs
@@ -1,3 +1,3 @@
 pub mod x86_64;
-
-pub mod ir;
\ No newline at end of file
+pub mod ir;
+mod tests;
\ No newline at end of file
diff --git a/yjit/src/backend/tests.rs b/yjit/src/backend/tests.rs
new file mode 100644
index 0000000000..45b8fdfb8a
--- /dev/null
+++ b/yjit/src/backend/tests.rs
@@ -0,0 +1,221 @@
+#![cfg(test)]
+
+use crate::asm::{CodeBlock};
+use crate::virtualmem::{CodePtr};
+use crate::backend::ir::*;
+use crate::cruby::*;
+use crate::core::*;
+use InsnOpnd::*;
+
+// Test that this function type checks
+fn gen_dup(
+    ctx: &mut Context,
+    asm: &mut Assembler,
+) {
+    let dup_val = ctx.ir_stack_pop(0);
+    let (mapping, tmp_type) = ctx.get_opnd_mapping(StackOpnd(0));
+
+    let loc0 = ctx.ir_stack_push_mapping((mapping, tmp_type));
+    asm.mov(loc0, dup_val);
+}
+
+fn guard_object_is_heap(
+    asm: &mut Assembler,
+    object_opnd: Opnd,
+    ctx: &mut Context,
+    side_exit: CodePtr,
+) {
+    asm.comment("guard object is heap");
+
+    // Test that the object is not an immediate
+    asm.test(object_opnd.clone(), Opnd::UImm(RUBY_IMMEDIATE_MASK as u64));
+    asm.jnz(Target::CodePtr(side_exit));
+
+    // Test that the object is not false or nil
+    asm.cmp(object_opnd.clone(), Opnd::UImm(Qnil.into()));
+    asm.jbe(Target::CodePtr(side_exit));
+}
+
+#[test]
+fn test_add() {
+    let mut asm = Assembler::new();
+    let out = asm.add(SP, Opnd::UImm(1));
+    asm.add(out, Opnd::UImm(2));
+}
+
+#[test]
+fn test_split_loads() {
+    let mut asm = Assembler::new();
+
+    let regs = Assembler::get_alloc_regs();
+
+    asm.add(
+        Opnd::mem(64, Opnd::Reg(regs[0]), 0),
+        Opnd::mem(64, Opnd::Reg(regs[1]), 0)
+    );
+
+    let result = asm.split_loads();
+    assert_eq!(result.insns.len(), 2);
+    assert_eq!(result.insns[0].op, Op::Load);
+}
+
+#[test]
+fn test_alloc_regs() {
+    let mut asm = Assembler::new();
+
+    // Get the first output that we're going to reuse later.
+    let out1 = asm.add(EC, Opnd::UImm(1));
+
+    // Pad some instructions in to make sure it can handle that.
+    asm.add(EC, Opnd::UImm(2));
+
+    // Get the second output we're going to reuse.
+    let out2 = asm.add(EC, Opnd::UImm(3));
+
+    // Pad another instruction.
+    asm.add(EC, Opnd::UImm(4));
+
+    // Reuse both the previously captured outputs.
+    asm.add(out1, out2);
+
+    // Now get a third output to make sure that the pool has registers to
+    // allocate now that the previous ones have been returned.
+    let out3 = asm.add(EC, Opnd::UImm(5));
+    asm.add(out3, Opnd::UImm(6));
+
+    // Here we're going to allocate the registers.
+    let result = asm.alloc_regs(Assembler::get_alloc_regs());
+
+    // Now we're going to verify that the out field has been appropriately
+    // updated for each of the instructions that needs it.
+    let regs = Assembler::get_alloc_regs();
+    assert_eq!(result.insns[0].out, Opnd::Reg(regs[0]));
+    assert_eq!(result.insns[2].out, Opnd::Reg(regs[1]));
+    assert_eq!(result.insns[5].out, Opnd::Reg(regs[0]));
+}
+
+fn setup_asm() -> (Assembler, CodeBlock) {
+    return (
+        Assembler::new(),
+        CodeBlock::new_dummy(1024)
+    );
+}
+
+// Test full codegen pipeline
+#[test]
+fn test_compile()
+{
+    let (mut asm, mut cb) = setup_asm();
+    let regs = Assembler::get_alloc_regs();
+
+    let out = asm.add(Opnd::Reg(regs[0]), Opnd::UImm(2));
+    asm.add(out, Opnd::UImm(2));
+
+    asm.compile_with_num_regs(&mut cb, 1);
+}
+
+// Test memory-to-memory move
+#[test]
+fn test_mov_mem2mem()
+{
+    let (mut asm, mut cb) = setup_asm();
+
+    asm.comment("check that comments work too");
+    asm.mov(Opnd::mem(64, SP, 0), Opnd::mem(64, SP, 8));
+
+    asm.compile_with_num_regs(&mut cb, 1);
+}
+
+// Test load of register into new register
+#[test]
+fn test_load_reg()
+{
+    let (mut asm, mut cb) = setup_asm();
+
+    let out = asm.load(SP);
+    asm.mov(Opnd::mem(64, SP, 0), out);
+
+    asm.compile_with_num_regs(&mut cb, 1);
+}
+
+// Multiple registers needed and register reuse
+#[test]
+fn test_reuse_reg()
+{
+    let (mut asm, mut cb) = setup_asm();
+
+    let v0 = asm.add(Opnd::mem(64, SP, 0), Opnd::UImm(1));
+    let v1 = asm.add(Opnd::mem(64, SP, 8), Opnd::UImm(1));
+    let v2 = asm.add(v0, Opnd::UImm(1));
+    asm.add(v0, v2);
+
+    asm.compile_with_num_regs(&mut cb, 2);
+}
+
+// 64-bit values can't be written directly to memory,
+// need to be split into one or more register movs first
+#[test]
+fn test_store_u64()
+{
+    let (mut asm, mut cb) = setup_asm();
+    asm.store(Opnd::mem(64, SP, 0), u64::MAX.into());
+
+    asm.compile_with_num_regs(&mut cb, 1);
+}
+
+// Use instruction output as base register for memory operand
+#[test]
+fn test_base_insn_out()
+{
+    let (mut asm, mut cb) = setup_asm();
+
+    // Load the pointer into a register
+    let ptr_reg = asm.load(Opnd::const_ptr(0 as *const u8));
+    let counter_opnd = Opnd::mem(64, ptr_reg, 0);
+
+    // Increment and store the updated value
+    asm.incr_counter(counter_opnd, 1.into() );
+
+    asm.compile_with_num_regs(&mut cb, 1);
+}
+
+#[test]
+fn test_c_call()
+{
+    extern "sysv64" fn dummy_c_fun(v0: usize, v1: usize)
+    {
+    }
+
+    let (mut asm, mut cb) = setup_asm();
+
+    asm.ccall(
+        dummy_c_fun as *const u8,
+        vec![Opnd::mem(64, SP, 0), Opnd::UImm(1)]
+    );
+
+    asm.compile_with_num_regs(&mut cb, 2);
+}
+
+#[test]
+fn test_lea_ret()
+{
+    let (mut asm, mut cb) = setup_asm();
+
+    let addr = asm.lea(Opnd::mem(64, SP, 0));
+    asm.cret(addr);
+
+    asm.compile_with_num_regs(&mut cb, 1);
+}
+
+#[test]
+fn test_jcc_label()
+{
+    let (mut asm, mut cb) = setup_asm();
+
+    let label = asm.new_label("foo");
+    asm.cmp(EC, EC);
+    asm.je(label);
+    asm.write_label(label);
+
+    asm.compile_with_num_regs(&mut cb, 1);
+}
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index 0c23781e20..4d24378370 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -67,7 +67,7 @@ impl From<Opnd> for X86Opnd {
 impl Assembler
 {
     /// Get the list of registers from which we can allocate on this platform
-    pub fn get_scratch_regs() -> Vec<Reg>
+    pub fn get_alloc_regs() -> Vec<Reg>
     {
         vec![
             RAX_REG,