Better offsets (#6315)

* Introduce InstructionOffset for AArch64

There are a lot of instructions on AArch64 where we take an offset
from PC in terms of the number of instructions. This is for loading
a value relative to the PC or for jumping.

We were usually accepting an A64Opnd or an i32. It can get
confusing and inconsistent though because sometimes you would
divide by 4 to get the number of instructions or multiply by 4 to
get the number of bytes.

This commit adds a struct that wraps an i32 in order to keep all of
that logic in one place. It makes it much easier to read and reason
about how these offsets are getting used.

* Use b instruction when the offset fits on AArch64

1 files changed, 47 insertions, 0 deletions

diff --git a/yjit/src/asm/arm64/arg/inst_offset.rs b/yjit/src/asm/arm64/arg/inst_offset.rs
new file mode 100644
index 0000000000..f4a6bc73a0
--- /dev/null
+++ b/yjit/src/asm/arm64/arg/inst_offset.rs
@@ -0,0 +1,47 @@
+/// There are a lot of instructions in the AArch64 architectrue that take an
+/// offset in terms of number of instructions. Usually they are jump
+/// instructions or instructions that load a value relative to the current PC.
+///
+/// This struct is used to mark those locations instead of a generic operand in
+/// order to give better clarity to the developer when reading the AArch64
+/// backend code. It also helps to clarify that everything is in terms of a
+/// number of instructions and not a number of bytes (i.e., the offset is the
+/// number of bytes divided by 4).
+#[derive(Copy, Clone)]
+pub struct InstructionOffset(i32);
+
+impl InstructionOffset {
+    /// Create a new instruction offset.
+    pub fn from_insns(insns: i32) -> Self {
+        InstructionOffset(insns)
+    }
+
+    /// Create a new instruction offset from a number of bytes.
+    pub fn from_bytes(bytes: i32) -> Self {
+        assert_eq!(bytes % 4, 0, "Byte offset must be a multiple of 4");
+        InstructionOffset(bytes / 4)
+    }
+}
+
+impl From<i32> for InstructionOffset {
+    /// Convert an i64 into an instruction offset.
+    fn from(value: i32) -> Self {
+        InstructionOffset(value)
+    }
+}
+
+impl From<InstructionOffset> for i32 {
+    /// Convert an instruction offset into a number of instructions as an i32.
+    fn from(offset: InstructionOffset) -> Self {
+        offset.0
+    }
+}
+
+impl From<InstructionOffset> for i64 {
+    /// Convert an instruction offset into a number of instructions as an i64.
+    /// This is useful for when we're checking how many bits this offset fits
+    /// into.
+    fn from(offset: InstructionOffset) -> Self {
+        offset.0.into()
+    }
+}