Arm64 progress (https://github.com/Shopify/ruby/pull/304)

* Get initial wiring up

* Split IncrCounter instruction

* Breakpoints in Arm64

* Support for ORR

* MOV instruction encodings

* Implement JmpOpnd and CRet

* Add ORN

* Add MVN

* PUSH, POP, CCALL for Arm64

* Some formatting and implement Op::Not for Arm64

* Consistent constants when working with the Arm64 SP

* Allow OR-ing values into the memory buffer

* Test lowering Arm64 ADD

* Emit unconditional jumps consistently in Arm64

* Begin emitting conditional jumps for A64

* Back out some labelref changes

* Remove label API that no longer exists

* Use a trait for the label encoders

* Encode nop

* Add in nops so jumps are the same width no matter what on Arm64

* Op::Jbe for CodePtr

* Pass src_addr and dst_addr instead of calculated offset to label refs

* Even more jump work for Arm64

* Fix up jumps to use consistent assertions

* Handle splitting Add, Sub, and Not insns for Arm64

* More Arm64 splits and various fixes

* PR feedback for Arm64 support

* Split up jumps and conditional jump logic

1 files changed, 19 insertions, 14 deletions

diff --git a/yjit/src/asm/mod.rs b/yjit/src/asm/mod.rs
index b54fc362b4..5723406aec 100644
--- a/yjit/src/asm/mod.rs
+++ b/yjit/src/asm/mod.rs
@@ -23,6 +23,14 @@ struct LabelRef {
 
     // Label which this refers to
     label_idx: usize,
+
+    /// The number of bytes that this label reference takes up in the memory.
+    /// It's necessary to know this ahead of time so that when we come back to
+    /// patch it it takes the same amount of space.
+    num_bytes: usize,
+
+    /// The object that knows how to encode the branch instruction.
+    encode: Box<dyn FnOnce(&mut CodeBlock, i64, i64)>
 }
 
 /// Block of memory into which instructions can be assembled
@@ -154,7 +162,7 @@ impl CodeBlock {
         self.get_ptr(self.write_pos)
     }
 
-    // Write a single byte at the current position
+    /// Write a single byte at the current position.
     pub fn write_byte(&mut self, byte: u8) {
         let write_ptr = self.get_write_ptr();
 
@@ -165,15 +173,15 @@ impl CodeBlock {
         }
     }
 
-    // Write multiple bytes starting from the current position
-    pub fn write_bytes(&mut self, bytes: &[u8]) {
+    /// Write multiple bytes starting from the current position.
+    fn write_bytes(&mut self, bytes: &[u8]) {
         for byte in bytes {
             self.write_byte(*byte);
         }
     }
 
-    // Write a signed integer over a given number of bits at the current position
-    pub fn write_int(&mut self, val: u64, num_bits: u32) {
+    /// Write an integer over the given number of bits at the current position.
+    fn write_int(&mut self, val: u64, num_bits: u32) {
         assert!(num_bits > 0);
         assert!(num_bits % 8 == 0);
 
@@ -219,14 +227,14 @@ impl CodeBlock {
     }
 
     // Add a label reference at the current write position
-    pub fn label_ref(&mut self, label_idx: usize) {
+    pub fn label_ref<E: 'static>(&mut self, label_idx: usize, num_bytes: usize, encode: E) where E: FnOnce(&mut CodeBlock, i64, i64) {
         assert!(label_idx < self.label_addrs.len());
 
         // Keep track of the reference
-        self.label_refs.push(LabelRef {
-            pos: self.write_pos,
-            label_idx,
-        });
+        self.label_refs.push(LabelRef { pos: self.write_pos, label_idx, num_bytes, encode: Box::new(encode) });
+
+        // Move past however many bytes the instruction takes up
+        self.write_pos += num_bytes;
     }
 
     // Link internal label references
@@ -242,11 +250,8 @@ impl CodeBlock {
             let label_addr = self.label_addrs[label_idx];
             assert!(label_addr < self.mem_size);
 
-            // Compute the offset from the reference's end to the label
-            let offset = (label_addr as i64) - ((ref_pos + 4) as i64);
-
             self.set_pos(ref_pos);
-            self.write_int(offset as u64, 32);
+            (label_ref.encode)(self, (ref_pos + label_ref.num_bytes) as i64, label_addr as i64);
         }
 
         self.write_pos = orig_pos;