Fix jumps (https://github.com/Shopify/ruby/pull/309)

* Jumps for A64 should be in # of instructions

* More splitting for Arm64

https://github.com/Shopify/ruby/pull/309

1 files changed, 47 insertions, 9 deletions

diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index 7e6a187f8f..0fee18c068 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -140,6 +140,14 @@ impl Assembler
 
                     asm.incr_counter(new_opnds[0], new_opnds[1]);
                 },
+                Op::JmpOpnd => {
+                    if let Opnd::Mem(_) = opnds[0] {
+                        let opnd0 = asm.load(opnds[0]);
+                        asm.jmp_opnd(opnd0);
+                    } else {
+                        asm.jmp_opnd(opnds[0]);
+                    }
+                },
                 Op::Mov => {
                     // The value that is being moved must be either a register
                     // or an immediate that can be encoded as a bitmask
@@ -150,7 +158,15 @@ impl Assembler
                         Opnd::Mem(_) | Opnd::Imm(_) => asm.load(opnds[1]),
                         Opnd::UImm(uimm) => {
                             if let Ok(encoded) = BitmaskImmediate::try_from(uimm) {
-                                opnds[1]
+                                if let Opnd::Mem(_) = opnds[0] {
+                                    // If the first operand is a memory operand,
+                                    // we're going to transform this into a
+                                    // store instruction, so we'll need to load
+                                    // this anyway.
+                                    asm.load(opnds[1])
+                                } else {
+                                    opnds[1]
+                                }
                             } else {
                                 asm.load(opnds[1])
                             }
@@ -158,9 +174,9 @@ impl Assembler
                         _ => unreachable!()
                     };
 
-                    /// If we're attempting to load into a memory operand, then
-                    /// we'll switch over to the store instruction. Otherwise
-                    /// we'll use the normal mov instruction.
+                    // If we're attempting to load into a memory operand, then
+                    // we'll switch over to the store instruction. Otherwise
+                    // we'll use the normal mov instruction.
                     match opnds[0] {
                         Opnd::Mem(_) => asm.store(opnds[0], value),
                         _ => asm.mov(opnds[0], value)
@@ -207,6 +223,26 @@ impl Assembler
     /// Returns a list of GC offsets
     pub fn arm64_emit(&mut self, cb: &mut CodeBlock) -> Vec<u32>
     {
+        /// Determine how many instructions it will take to represent moving
+        /// this value into a register. Note that the return value of this
+        /// function must correspond to how many instructions are used to
+        /// represent this load in the emit_load_value function.
+        fn emit_load_size(value: u64) -> u8 {
+            if BitmaskImmediate::try_from(value).is_ok() {
+                return 1;
+            }
+
+            if value < (1 << 16) {
+                1
+            } else if value < (1 << 32) {
+                2
+            } else if value < (1 << 48) {
+                3
+            } else {
+                4
+            }
+        }
+
         /// Emit the required instructions to load the given value into the
         /// given register. Our goal here is to use as few instructions as
         /// possible to get this value into the register.
@@ -275,7 +311,7 @@ impl Assembler
                             // If we get to this instruction, then the condition
                             // wasn't met, in which case we'll jump past the
                             // next instruction that performs the direct jump.
-                            b(cb, A64Opnd::new_imm(4));
+                            b(cb, A64Opnd::new_imm(1));
 
                             // Here we'll perform the direct jump to the target.
                             b(cb, A64Opnd::new_imm(offset / 4));
@@ -283,8 +319,10 @@ impl Assembler
                             // If we get to this instruction, then the condition
                             // wasn't met, in which case we'll jump past the
                             // next instruction that perform the direct jump.
-                            b(cb, A64Opnd::new_imm(8));
-                            emit_load_value(cb, X29, dst_addr as u64);
+                            let value = dst_addr as u64;
+
+                            b(cb, A64Opnd::new_imm(emit_load_size(value).into()));
+                            emit_load_value(cb, X29, value);
                             br(cb, X29);
                         }
                     }
@@ -392,10 +430,10 @@ impl Assembler
                 },
                 Op::CPush => {
                     add(cb, C_SP_REG, C_SP_REG, C_SP_STEP);
-                    mov(cb, A64Opnd::new_mem(64, C_SP_REG, 0), insn.opnds[0].into());
+                    stur(cb, insn.opnds[0].into(), A64Opnd::new_mem(64, C_SP_REG, 0));
                 },
                 Op::CPop => {
-                    mov(cb, insn.out.into(), A64Opnd::new_mem(64, C_SP_REG, 0));
+                    ldur(cb, insn.opnds[0].into(), A64Opnd::new_mem(64, C_SP_REG, 0));
                     sub(cb, C_SP_REG, C_SP_REG, C_SP_STEP);
                 },
                 Op::CCall => {