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

8 files changed, 474 insertions, 43 deletions

diff --git a/yjit/src/asm/arm64/inst/breakpoint.rs b/yjit/src/asm/arm64/inst/breakpoint.rs
new file mode 100644
index 0000000000..be4920ac76
--- /dev/null
+++ b/yjit/src/asm/arm64/inst/breakpoint.rs
@@ -0,0 +1,55 @@
+/// The struct that represents an A64 breakpoint instruction that can be encoded.
+///
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+/// | 31 30 29 28 | 27 26 25 24 | 23 22 21 20 | 19 18 17 16 | 15 14 13 12 | 11 10 09 08 | 07 06 05 04 | 03 02 01 00 |
+/// |  1  1  0  1    0  1  0  0    0  0  1                                                          0    0  0  0  0 |
+/// |                                      imm16..................................................                  |
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+///
+pub struct Breakpoint {
+    /// The value to be captured by ESR_ELx.ISS
+    imm16: u16
+}
+
+impl Breakpoint {
+    /// BRK
+    /// https://developer.arm.com/documentation/ddi0596/2020-12/Base-Instructions/BRK--Breakpoint-instruction-
+    pub fn brk(imm16: u16) -> Self {
+        Self { imm16 }
+    }
+}
+
+/// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/Branches--Exception-Generating-and-System-instructions?lang=en#control
+const FAMILY: u32 = 0b101;
+
+impl From<Breakpoint> for u32 {
+    /// Convert an instruction into a 32-bit value.
+    fn from(inst: Breakpoint) -> Self {
+        let imm16 = inst.imm16 as u32;
+
+        0
+        | (0b11 << 30)
+        | (FAMILY << 26)
+        | (1 << 21)
+        | (imm16 << 5)
+    }
+}
+
+impl From<Breakpoint> for [u8; 4] {
+    /// Convert an instruction into a 4 byte array.
+    fn from(inst: Breakpoint) -> [u8; 4] {
+        let result: u32 = inst.into();
+        result.to_le_bytes()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_brk() {
+        let result: u32 = Breakpoint::brk(7).into();
+        assert_eq!(0xd42000e0, result);
+    }
+}
diff --git a/yjit/src/asm/arm64/inst/call.rs b/yjit/src/asm/arm64/inst/call.rs
index 6f23acf9f5..8d65359f77 100644
--- a/yjit/src/asm/arm64/inst/call.rs
+++ b/yjit/src/asm/arm64/inst/call.rs
@@ -1,22 +1,41 @@
-/// The struct that represents an A64 branch with link instruction that can be
-/// encoded.
+/// The operation to perform for this instruction.
+enum Op {
+    /// Branch directly, with a hint that this is not a subroutine call or
+    /// return.
+    Branch = 0,
+
+    /// Branch directly, with a hint that this is a subroutine call or return.
+    BranchWithLink = 1
+}
+
+/// The struct that represents an A64 branch with our without link instruction
+/// that can be encoded.
 ///
 /// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
 /// | 31 30 29 28 | 27 26 25 24 | 23 22 21 20 | 19 18 17 16 | 15 14 13 12 | 11 10 09 08 | 07 06 05 04 | 03 02 01 00 |
-/// |  1  0  0  1    0  1                                                                                           |
-/// |                     imm26.................................................................................... |
+/// |     0  0  1    0  1                                                                                           |
+/// | op                  imm26.................................................................................... |
 /// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
 ///
 pub struct Call {
     /// The PC-relative offset to jump to (which will be multiplied by 4).
-    imm26: i32
+    imm26: i32,
+
+    /// The operation to perform for this instruction.
+    op: Op
 }
 
 impl Call {
+    /// B
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/B--Branch-
+    pub fn b(imm26: i32) -> Self {
+        Self { imm26, op: Op::Branch }
+    }
+
     /// BL
     /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/BL--Branch-with-Link-?lang=en
     pub fn bl(imm26: i32) -> Self {
-        Self { imm26 }
+        Self { imm26, op: Op::BranchWithLink }
     }
 }
 
@@ -29,7 +48,7 @@ impl From<Call> for u32 {
         let imm26 = (inst.imm26 as u32) & ((1 << 26) - 1);
 
         0
-        | (1 << 31)
+        | ((inst.op as u32) << 31)
         | (FAMILY << 26)
         | imm26
     }
@@ -64,4 +83,22 @@ mod tests {
         let result: u32 = Call::bl(-256).into();
         assert_eq!(0x97ffff00, result);
     }
+
+    #[test]
+    fn test_b() {
+        let result: u32 = Call::b(0).into();
+        assert_eq!(0x14000000, result);
+    }
+
+    #[test]
+    fn test_b_positive() {
+        let result: u32 = Call::b(256).into();
+        assert_eq!(0x14000100, result);
+    }
+
+    #[test]
+    fn test_b_negative() {
+        let result: u32 = Call::b(-256).into();
+        assert_eq!(0x17ffff00, result);
+    }
 }
diff --git a/yjit/src/asm/arm64/inst/logical_imm.rs b/yjit/src/asm/arm64/inst/logical_imm.rs
index cc2a16cbdc..13865697f6 100644
--- a/yjit/src/asm/arm64/inst/logical_imm.rs
+++ b/yjit/src/asm/arm64/inst/logical_imm.rs
@@ -5,6 +5,9 @@ enum Opc {
     /// The AND operation.
     And = 0b00,
 
+    /// The ORR operation.
+    Orr = 0b01,
+
     /// The ANDS operation.
     Ands = 0b11
 }
@@ -12,7 +15,7 @@ enum Opc {
 /// The struct that represents an A64 bitwise immediate instruction that can be
 /// encoded.
 ///
-/// AND/ANDS (immediate)
+/// AND/ORR/ANDS (immediate)
 /// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
 /// | 31 30 29 28 | 27 26 25 24 | 23 22 21 20 | 19 18 17 16 | 15 14 13 12 | 11 10 09 08 | 07 06 05 04 | 03 02 01 00 |
 /// |           1    0  0  1  0    0                                                                                |
@@ -37,19 +40,31 @@ pub struct LogicalImm {
 }
 
 impl LogicalImm {
-    /// AND (immediate)
+    /// AND (bitmask immediate)
     /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/AND--immediate---Bitwise-AND--immediate--?lang=en
     pub fn and(rd: u8, rn: u8, imm: BitmaskImmediate, num_bits: u8) -> Self {
         Self { rd, rn, imm, opc: Opc::And, sf: num_bits.into() }
     }
 
-    /// ANDS (immediate)
+    /// ANDS (bitmask immediate)
     /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/ANDS--immediate---Bitwise-AND--immediate---setting-flags-?lang=en
     pub fn ands(rd: u8, rn: u8, imm: BitmaskImmediate, num_bits: u8) -> Self {
         Self { rd, rn, imm, opc: Opc::Ands, sf: num_bits.into() }
     }
 
-    /// TST (immediate)
+    /// MOV (bitmask immediate)
+    /// https://developer.arm.com/documentation/ddi0596/2020-12/Base-Instructions/MOV--bitmask-immediate---Move--bitmask-immediate---an-alias-of-ORR--immediate--?lang=en
+    pub fn mov(rd: u8, imm: BitmaskImmediate, num_bits: u8) -> Self {
+        Self { rd, rn: 0b11111, imm, opc: Opc::Orr, sf: num_bits.into() }
+    }
+
+    /// ORR (bitmask immediate)
+    /// https://developer.arm.com/documentation/ddi0596/2020-12/Base-Instructions/ORR--immediate---Bitwise-OR--immediate--
+    pub fn orr(rd: u8, rn: u8, imm: BitmaskImmediate, num_bits: u8) -> Self {
+        Self { rd, rn, imm, opc: Opc::Orr, sf: num_bits.into() }
+    }
+
+    /// TST (bitmask immediate)
     /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/TST--immediate---Test-bits--immediate---an-alias-of-ANDS--immediate--?lang=en
     pub fn tst(rn: u8, imm: BitmaskImmediate, num_bits: u8) -> Self {
         Self::ands(31, rn, imm, num_bits)
@@ -101,6 +116,20 @@ mod tests {
     }
 
     #[test]
+    fn test_mov() {
+        let inst = LogicalImm::mov(0, 0x5555555555555555.try_into().unwrap(), 64);
+        let result: u32 = inst.into();
+        assert_eq!(0xb200f3e0, result);
+    }
+
+    #[test]
+    fn test_orr() {
+        let inst = LogicalImm::orr(0, 1, 7.try_into().unwrap(), 64);
+        let result: u32 = inst.into();
+        assert_eq!(0xb2400820, result);
+    }
+
+    #[test]
     fn test_tst() {
         let inst = LogicalImm::tst(1, 7.try_into().unwrap(), 64);
         let result: u32 = inst.into();
diff --git a/yjit/src/asm/arm64/inst/logical_reg.rs b/yjit/src/asm/arm64/inst/logical_reg.rs
index 3feb3350ab..5d7954c587 100644
--- a/yjit/src/asm/arm64/inst/logical_reg.rs
+++ b/yjit/src/asm/arm64/inst/logical_reg.rs
@@ -1,5 +1,14 @@
 use super::super::arg::Sf;
 
+/// Whether or not this is a NOT instruction.
+enum N {
+    /// This is not a NOT instruction.
+    No = 0,
+
+    /// This is a NOT instruction.
+    Yes = 1
+}
+
 /// The type of shift to perform on the second operand register.
 enum Shift {
     LSL = 0b00, // logical shift left (unsigned)
@@ -13,6 +22,9 @@ enum Opc {
     /// The AND operation.
     And = 0b00,
 
+    /// The ORR operation.
+    Orr = 0b01,
+
     /// The ANDS operation.
     Ands = 0b11
 }
@@ -20,11 +32,11 @@ enum Opc {
 /// The struct that represents an A64 logical register instruction that can be
 /// encoded.
 ///
-/// AND/ANDS (shifted register)
+/// AND/ORR/ANDS (shifted register)
 /// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
 /// | 31 30 29 28 | 27 26 25 24 | 23 22 21 20 | 19 18 17 16 | 15 14 13 12 | 11 10 09 08 | 07 06 05 04 | 03 02 01 00 |
-/// |           0    1  0  1  0          0                                                                          |
-/// | sf opc..                    shift    rm..............   imm6............... rn.............. rd.............. |
+/// |           0    1  0  1  0                                                                                     |
+/// | sf opc..                    shift N  rm..............   imm6............... rn.............. rd.............. |
 /// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
 ///
 pub struct LogicalReg {
@@ -40,6 +52,9 @@ pub struct LogicalReg {
     /// The register number of the second operand register.
     rm: u8,
 
+    /// Whether or not this is a NOT instruction.
+    n: N,
+
     /// The type of shift to perform on the second operand register.
     shift: Shift,
 
@@ -54,19 +69,43 @@ impl LogicalReg {
     /// AND (shifted register)
     /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/AND--shifted-register---Bitwise-AND--shifted-register--?lang=en
     pub fn and(rd: u8, rn: u8, rm: u8, num_bits: u8) -> Self {
-        Self { rd, rn, imm6: 0, rm, shift: Shift::LSL, opc: Opc::And, sf: num_bits.into() }
+        Self { rd, rn, imm6: 0, rm, n: N::No, shift: Shift::LSL, opc: Opc::And, sf: num_bits.into() }
     }
 
     /// ANDS (shifted register)
     /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/ANDS--shifted-register---Bitwise-AND--shifted-register---setting-flags-?lang=en
     pub fn ands(rd: u8, rn: u8, rm: u8, num_bits: u8) -> Self {
-        Self { rd, rn, imm6: 0, rm, shift: Shift::LSL, opc: Opc::Ands, sf: num_bits.into() }
+        Self { rd, rn, imm6: 0, rm, n: N::No, shift: Shift::LSL, opc: Opc::Ands, sf: num_bits.into() }
+    }
+
+    /// MOV (register)
+    /// https://developer.arm.com/documentation/ddi0596/2020-12/Base-Instructions/MOV--register---Move--register---an-alias-of-ORR--shifted-register--?lang=en
+    pub fn mov(rd: u8, rm: u8, num_bits: u8) -> Self {
+        Self { rd, rn: 0b11111, imm6: 0, rm, n: N::No, shift: Shift::LSL, opc: Opc::Orr, sf: num_bits.into() }
+    }
+
+    /// MVN (shifted register)
+    /// https://developer.arm.com/documentation/ddi0596/2020-12/Base-Instructions/MVN--Bitwise-NOT--an-alias-of-ORN--shifted-register--?lang=en
+    pub fn mvn(rd: u8, rm: u8, num_bits: u8) -> Self {
+        Self { rd, rn: 0b11111, imm6: 0, rm, n: N::Yes, shift: Shift::LSL, opc: Opc::Orr, sf: num_bits.into() }
+    }
+
+    /// ORN (shifted register)
+    /// https://developer.arm.com/documentation/ddi0596/2020-12/Base-Instructions/ORN--shifted-register---Bitwise-OR-NOT--shifted-register--
+    pub fn orn(rd: u8, rn: u8, rm: u8, num_bits: u8) -> Self {
+        Self { rd, rn, imm6: 0, rm, n: N::Yes, shift: Shift::LSL, opc: Opc::Orr, sf: num_bits.into() }
+    }
+
+    /// ORR (shifted register)
+    /// https://developer.arm.com/documentation/ddi0596/2020-12/Base-Instructions/ORR--shifted-register---Bitwise-OR--shifted-register--
+    pub fn orr(rd: u8, rn: u8, rm: u8, num_bits: u8) -> Self {
+        Self { rd, rn, imm6: 0, rm, n: N::No, shift: Shift::LSL, opc: Opc::Orr, sf: num_bits.into() }
     }
 
     /// TST (shifted register)
     /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/TST--shifted-register---Test--shifted-register---an-alias-of-ANDS--shifted-register--?lang=en
     pub fn tst(rn: u8, rm: u8, num_bits: u8) -> Self {
-        Self { rd: 31, rn, imm6: 0, rm, shift: Shift::LSL, opc: Opc::Ands, sf: num_bits.into() }
+        Self { rd: 31, rn, imm6: 0, rm, n: N::No, shift: Shift::LSL, opc: Opc::Ands, sf: num_bits.into() }
     }
 }
 
@@ -83,6 +122,7 @@ impl From<LogicalReg> for u32 {
         | ((inst.opc as u32) << 29)
         | (FAMILY << 25)
         | ((inst.shift as u32) << 22)
+        | ((inst.n as u32) << 21)
         | ((inst.rm as u32) << 16)
         | (imm6 << 10)
         | ((inst.rn as u32) << 5)
@@ -117,6 +157,34 @@ mod tests {
     }
 
     #[test]
+    fn test_mov() {
+        let inst = LogicalReg::mov(0, 1, 64);
+        let result: u32 = inst.into();
+        assert_eq!(0xaa0103e0, result);
+    }
+
+    #[test]
+    fn test_mvn() {
+        let inst = LogicalReg::mvn(0, 1, 64);
+        let result: u32 = inst.into();
+        assert_eq!(0xaa2103e0, result);
+    }
+
+    #[test]
+    fn test_orn() {
+        let inst = LogicalReg::orn(0, 1, 2, 64);
+        let result: u32 = inst.into();
+        assert_eq!(0xaa220020, result);
+    }
+
+    #[test]
+    fn test_orr() {
+        let inst = LogicalReg::orr(0, 1, 2, 64);
+        let result: u32 = inst.into();
+        assert_eq!(0xaa020020, result);
+    }
+
+    #[test]
     fn test_tst() {
         let inst = LogicalReg::tst(0, 1, 64);
         let result: u32 = inst.into();
diff --git a/yjit/src/asm/arm64/inst/mod.rs b/yjit/src/asm/arm64/inst/mod.rs
index 2f0e708999..ae589ca564 100644
--- a/yjit/src/asm/arm64/inst/mod.rs
+++ b/yjit/src/asm/arm64/inst/mod.rs
@@ -4,6 +4,7 @@
 mod atomic;
 mod branch;
 mod branch_cond;
+mod breakpoint;
 mod call;
 mod data_imm;
 mod data_reg;
@@ -11,12 +12,14 @@ mod load;
 mod logical_imm;
 mod logical_reg;
 mod mov;
+mod nop;
 mod shift_imm;
 mod store;
 
 pub use atomic::Atomic;
 pub use branch::Branch;
 pub use branch_cond::BranchCond;
+pub use breakpoint::Breakpoint;
 pub use call::Call;
 pub use data_imm::DataImm;
 pub use data_reg::DataReg;
@@ -24,5 +27,6 @@ pub use load::Load;
 pub use logical_imm::LogicalImm;
 pub use logical_reg::LogicalReg;
 pub use mov::Mov;
+pub use nop::Nop;
 pub use shift_imm::ShiftImm;
 pub use store::Store;
diff --git a/yjit/src/asm/arm64/inst/nop.rs b/yjit/src/asm/arm64/inst/nop.rs
new file mode 100644
index 0000000000..a99f8d34b7
--- /dev/null
+++ b/yjit/src/asm/arm64/inst/nop.rs
@@ -0,0 +1,44 @@
+/// The struct that represents an A64 nop instruction that can be encoded.
+///
+/// NOP
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+/// | 31 30 29 28 | 27 26 25 24 | 23 22 21 20 | 19 18 17 16 | 15 14 13 12 | 11 10 09 08 | 07 06 05 04 | 03 02 01 00 |
+/// |  1  1  0  1    0  1  0  1    0  0  0  0    0  0  1  1    0  0  1  0    0  0  0  0    0  0  0  1    1  1  1  1 |
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+///
+pub struct Nop;
+
+impl Nop {
+    /// NOP
+    /// https://developer.arm.com/documentation/ddi0602/2022-03/Base-Instructions/NOP--No-Operation-
+    pub fn nop() -> Self {
+        Self {}
+    }
+}
+
+impl From<Nop> for u32 {
+    /// Convert an instruction into a 32-bit value.
+    fn from(inst: Nop) -> Self {
+        0b11010101000000110010000000011111
+    }
+}
+
+impl From<Nop> for [u8; 4] {
+    /// Convert an instruction into a 4 byte array.
+    fn from(inst: Nop) -> [u8; 4] {
+        let result: u32 = inst.into();
+        result.to_le_bytes()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_nop() {
+        let inst = Nop::nop();
+        let result: u32 = inst.into();
+        assert_eq!(0xd503201f, result);
+    }
+}
diff --git a/yjit/src/asm/arm64/mod.rs b/yjit/src/asm/arm64/mod.rs
index 24f349d589..ced8b262c5 100644
--- a/yjit/src/asm/arm64/mod.rs
+++ b/yjit/src/asm/arm64/mod.rs
@@ -6,12 +6,15 @@ mod arg;
 mod inst;
 mod opnd;
 
-use arg::*;
 use inst::*;
-use opnd::*;
+
+// We're going to make these public to make using these things easier in the
+// backend (so they don't have to have knowledge about the submodule).
+pub use arg::*;
+pub use opnd::*;
 
 /// Checks that a signed value fits within the specified number of bits.
-const fn imm_fits_bits(imm: i64, num_bits: u8) -> bool {
+pub const fn imm_fits_bits(imm: i64, num_bits: u8) -> bool {
     let minimum = if num_bits == 64 { i64::MIN } else { -2_i64.pow((num_bits as u32) - 1) };
     let maximum = if num_bits == 64 { i64::MAX } else { 2_i64.pow((num_bits as u32) - 1) - 1 };
 
@@ -19,7 +22,7 @@ const fn imm_fits_bits(imm: i64, num_bits: u8) -> bool {
 }
 
 /// Checks that an unsigned value fits within the specified number of bits.
-const fn uimm_fits_bits(uimm: u64, num_bits: u8) -> bool {
+pub const fn uimm_fits_bits(uimm: u64, num_bits: u8) -> bool {
     let maximum = if num_bits == 64 { u64::MAX } else { 2_u64.pow(num_bits as u32) - 1 };
 
     uimm <= maximum
@@ -115,12 +118,39 @@ pub fn ands(cb: &mut CodeBlock, rd: A64Opnd, rn: A64Opnd, rm: A64Opnd) {
     cb.write_bytes(&bytes);
 }
 
+/// Whether or not the offset between two instructions fits into the branch with
+/// or without link instruction. If it doesn't, then we have to load the value
+/// into a register first.
+pub const fn b_offset_fits_bits(offset: i64) -> bool {
+    imm_fits_bits(offset, 26)
+}
+
+/// B - branch without link (offset is number of instructions to jump)
+pub fn b(cb: &mut CodeBlock, imm26: A64Opnd) {
+    let bytes: [u8; 4] = match imm26 {
+        A64Opnd::Imm(imm26) => {
+            assert!(b_offset_fits_bits(imm26), "The immediate operand must be 26 bits or less.");
+
+            Call::b(imm26 as i32).into()
+        },
+        _ => panic!("Invalid operand combination to b instruction.")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
+/// Whether or not the offset between two instructions fits into the b.cond
+/// instruction. If it doesn't, then we have to load the value into a register
+/// first, then use the b.cond instruction to skip past a direct jump.
+pub const fn bcond_offset_fits_bits(offset: i64) -> bool {
+    imm_fits_bits(offset, 21) && (offset & 0b11 == 0)
+}
+
 /// B.cond - branch to target if condition is true
 pub fn bcond(cb: &mut CodeBlock, cond: Condition, byte_offset: A64Opnd) {
     let bytes: [u8; 4] = match byte_offset {
         A64Opnd::Imm(imm) => {
-            assert!(imm_fits_bits(imm, 21), "The immediate operand must be 21 bits or less.");
-            assert!(imm & 0b11 == 0, "The immediate operand must be aligned to a 2-bit boundary.");
+            assert!(bcond_offset_fits_bits(imm), "The immediate operand must be 21 bits or less and be aligned to a 2-bit boundary.");
 
             BranchCond::bcond(cond, imm as i32).into()
         },
@@ -134,7 +164,7 @@ pub fn bcond(cb: &mut CodeBlock, cond: Condition, byte_offset: A64Opnd) {
 pub fn bl(cb: &mut CodeBlock, imm26: A64Opnd) {
     let bytes: [u8; 4] = match imm26 {
         A64Opnd::Imm(imm26) => {
-            assert!(imm_fits_bits(imm26, 26), "The immediate operand must be 26 bits or less.");
+            assert!(b_offset_fits_bits(imm26), "The immediate operand must be 26 bits or less.");
 
             Call::bl(imm26 as i32).into()
         },
@@ -154,6 +184,20 @@ pub fn br(cb: &mut CodeBlock, rn: A64Opnd) {
     cb.write_bytes(&bytes);
 }
 
+/// BRK - create a breakpoint
+pub fn brk(cb: &mut CodeBlock, imm16: A64Opnd) {
+    let bytes: [u8; 4] = match imm16 {
+        A64Opnd::None => Breakpoint::brk(0).into(),
+        A64Opnd::UImm(imm16) => {
+            assert!(uimm_fits_bits(imm16, 16), "The immediate operand must be 16 bits or less.");
+            Breakpoint::brk(imm16 as u16).into()
+        },
+        _ => panic!("Invalid operand combination to brk instruction.")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
 /// CMP - compare rn and rm, update flags
 pub fn cmp(cb: &mut CodeBlock, rn: A64Opnd, rm: A64Opnd) {
     let bytes: [u8; 4] = match (rn, rm) {
@@ -196,6 +240,11 @@ pub fn ldaddal(cb: &mut CodeBlock, rs: A64Opnd, rt: A64Opnd, rn: A64Opnd) {
 /// LDUR - load a memory address into a register
 pub fn ldur(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
     let bytes: [u8; 4] = match (rt, rn) {
+        (A64Opnd::Reg(rt), A64Opnd::Reg(rn)) => {
+            assert!(rt.num_bits == rn.num_bits, "All operands must be of the same size.");
+
+            Load::ldur(rt.reg_no, rn.reg_no, 0, rt.num_bits).into()
+        },
         (A64Opnd::Reg(rt), A64Opnd::Mem(rn)) => {
             assert!(rt.num_bits == rn.num_bits, "Expected registers to be the same size");
             assert!(imm_fits_bits(rn.disp.into(), 9), "Expected displacement to be 9 bits or less");
@@ -238,6 +287,23 @@ pub fn lsr(cb: &mut CodeBlock, rd: A64Opnd, rn: A64Opnd, shift: A64Opnd) {
     cb.write_bytes(&bytes);
 }
 
+/// MOV - move a value in a register to another register
+pub fn mov(cb: &mut CodeBlock, rd: A64Opnd, rm: A64Opnd) {
+    let bytes: [u8; 4] = match (rd, rm) {
+        (A64Opnd::Reg(rd), A64Opnd::Reg(rm)) => {
+            assert!(rd.num_bits == rm.num_bits, "Expected registers to be the same size");
+
+            LogicalReg::mov(rd.reg_no, rm.reg_no, rd.num_bits).into()
+        },
+        (A64Opnd::Reg(rd), A64Opnd::UImm(imm)) => {
+            LogicalImm::mov(rd.reg_no, imm.try_into().unwrap(), rd.num_bits).into()
+        },
+        _ => panic!("Invalid operand combination to mov instruction")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
 /// MOVK - move a 16 bit immediate into a register, keep the other bits in place
 pub fn movk(cb: &mut CodeBlock, rd: A64Opnd, imm16: A64Opnd, shift: u8) {
     let bytes: [u8; 4] = match (rd, imm16) {
@@ -266,6 +332,63 @@ pub fn movz(cb: &mut CodeBlock, rd: A64Opnd, imm16: A64Opnd, shift: u8) {
     cb.write_bytes(&bytes);
 }
 
+/// MVN - move a value in a register to another register, negating it
+pub fn mvn(cb: &mut CodeBlock, rd: A64Opnd, rm: A64Opnd) {
+    let bytes: [u8; 4] = match (rd, rm) {
+        (A64Opnd::Reg(rd), A64Opnd::Reg(rm)) => {
+            assert!(rd.num_bits == rm.num_bits, "Expected registers to be the same size");
+
+            LogicalReg::mvn(rd.reg_no, rm.reg_no, rd.num_bits).into()
+        },
+        _ => panic!("Invalid operand combination to mvn instruction")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
+/// NOP - no-operation, used for alignment purposes
+pub fn nop(cb: &mut CodeBlock) {
+    let bytes: [u8; 4] = Nop::nop().into();
+
+    cb.write_bytes(&bytes);
+}
+
+/// ORN - perform a bitwise OR of rn and NOT rm, put the result in rd, don't update flags
+pub fn orn(cb: &mut CodeBlock, rd: A64Opnd, rn: A64Opnd, rm: A64Opnd) {
+    let bytes: [u8; 4] = match (rd, rn, rm) {
+        (A64Opnd::Reg(rd), A64Opnd::Reg(rn), A64Opnd::Reg(rm)) => {
+            assert!(rd.num_bits == rn.num_bits && rn.num_bits == rm.num_bits, "Expected registers to be the same size");
+
+            LogicalReg::orn(rd.reg_no, rn.reg_no, rm.reg_no, rd.num_bits).into()
+        },
+        _ => panic!("Invalid operand combination to orn instruction.")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
+/// ORR - perform a bitwise OR of rn and rm, put the result in rd, don't update flags
+pub fn orr(cb: &mut CodeBlock, rd: A64Opnd, rn: A64Opnd, rm: A64Opnd) {
+    let bytes: [u8; 4] = match (rd, rn, rm) {
+        (A64Opnd::Reg(rd), A64Opnd::Reg(rn), A64Opnd::Reg(rm)) => {
+            assert!(
+                rd.num_bits == rn.num_bits && rn.num_bits == rm.num_bits,
+                "All operands must be of the same size."
+            );
+
+            LogicalReg::orr(rd.reg_no, rn.reg_no, rm.reg_no, rd.num_bits).into()
+        },
+        (A64Opnd::Reg(rd), A64Opnd::Reg(rn), A64Opnd::UImm(imm)) => {
+            assert!(rd.num_bits == rn.num_bits, "rd and rn must be of the same size.");
+
+            LogicalImm::orr(rd.reg_no, rn.reg_no, imm.try_into().unwrap(), rd.num_bits).into()
+        },
+        _ => panic!("Invalid operand combination to orr instruction."),
+    };
+
+    cb.write_bytes(&bytes);
+}
+
 /// STUR - store a value in a register at a memory address
 pub fn stur(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
     let bytes: [u8; 4] = match (rt, rn) {
@@ -435,6 +558,11 @@ mod tests {
     }
 
     #[test]
+    fn test_b() {
+        check_bytes("00040014", |cb| b(cb, A64Opnd::new_imm(1024)));
+    }
+
+    #[test]
     fn test_bl() {
         check_bytes("00040094", |cb| bl(cb, A64Opnd::new_imm(1024)));
     }
@@ -445,6 +573,16 @@ mod tests {
     }
 
     #[test]
+    fn test_brk_none() {
+        check_bytes("000020d4", |cb| brk(cb, A64Opnd::None));
+    }
+
+    #[test]
+    fn test_brk_uimm() {
+        check_bytes("c00120d4", |cb| brk(cb, A64Opnd::new_uimm(14)));
+    }
+
+    #[test]
     fn test_cmp_register() {
         check_bytes("5f010beb", |cb| cmp(cb, X10, X11));
     }
@@ -460,8 +598,13 @@ mod tests {
     }
 
     #[test]
-    fn test_ldur() {
-        check_bytes("20b047f8", |cb| ldur(cb, X0, A64Opnd::new_mem(X1, 123)));
+    fn test_ldur_memory() {
+        check_bytes("20b047f8", |cb| ldur(cb, X0, A64Opnd::new_mem(64, X1, 123)));
+    }
+
+    #[test]
+    fn test_ldur_register() {
+        check_bytes("200040f8", |cb| ldur(cb, X0, X1));
     }
 
     #[test]
@@ -475,6 +618,16 @@ mod tests {
     }
 
     #[test]
+    fn test_mov_registers() {
+        check_bytes("ea030baa", |cb| mov(cb, X10, X11));
+    }
+
+    #[test]
+    fn test_mov_immediate() {
+        check_bytes("eaf300b2", |cb| mov(cb, X10, A64Opnd::new_uimm(0x5555555555555555)));
+    }
+
+    #[test]
     fn test_movk() {
         check_bytes("600fa0f2", |cb| movk(cb, X0, A64Opnd::new_uimm(123), 16));
     }
@@ -485,6 +638,31 @@ mod tests {
     }
 
     #[test]
+    fn test_mvn() {
+        check_bytes("ea032baa", |cb| mvn(cb, X10, X11));
+    }
+
+    #[test]
+    fn test_nop() {
+        check_bytes("1f2003d5", |cb| nop(cb));
+    }
+
+    #[test]
+    fn test_orn() {
+        check_bytes("6a012caa", |cb| orn(cb, X10, X11, X12));
+    }
+
+    #[test]
+    fn test_orr_register() {
+        check_bytes("6a010caa", |cb| orr(cb, X10, X11, X12));
+    }
+
+    #[test]
+    fn test_orr_immediate() {
+        check_bytes("6a0940b2", |cb| orr(cb, X10, X11, A64Opnd::new_uimm(7)));
+    }
+
+    #[test]
     fn test_ret_none() {
         check_bytes("c0035fd6", |cb| ret(cb, A64Opnd::None));
     }
@@ -496,7 +674,7 @@ mod tests {
 
     #[test]
     fn test_stur() {
-        check_bytes("6a0108f8", |cb| stur(cb, X10, A64Opnd::new_mem(X11, 128)));
+        check_bytes("6a0108f8", |cb| stur(cb, X10, A64Opnd::new_mem(64, X11, 128)));
     }
 
     #[test]
diff --git a/yjit/src/asm/arm64/opnd.rs b/yjit/src/asm/arm64/opnd.rs
index aa73d438fe..6c06d2db3c 100644
--- a/yjit/src/asm/arm64/opnd.rs
+++ b/yjit/src/asm/arm64/opnd.rs
@@ -11,6 +11,15 @@ pub struct A64Reg
     pub reg_no: u8,
 }
 
+impl A64Reg {
+    pub fn sub_reg(&self, num_bits: u8) -> Self {
+        assert!(num_bits == 32 || num_bits == 64);
+        assert!(num_bits <= self.num_bits);
+
+        Self { num_bits, reg_no: self.reg_no }
+    }
+}
+
 #[derive(Clone, Copy, Debug)]
 pub struct A64Mem
 {
@@ -25,14 +34,10 @@ pub struct A64Mem
 }
 
 impl A64Mem {
-    pub fn new(reg: A64Opnd, disp: i32) -> Self {
+    pub fn new(num_bits: u8, reg: A64Opnd, disp: i32) -> Self {
         match reg {
             A64Opnd::Reg(reg) => {
-                Self {
-                    num_bits: reg.num_bits,
-                    base_reg_no: reg.reg_no,
-                    disp
-                }
+                Self { num_bits, base_reg_no: reg.reg_no, disp }
             },
             _ => panic!("Expected register operand")
         }
@@ -70,8 +75,8 @@ impl A64Opnd {
     }
 
     /// Creates a new memory operand.
-    pub fn new_mem(reg: A64Opnd, disp: i32) -> Self {
-        A64Opnd::Mem(A64Mem::new(reg, disp))
+    pub fn new_mem(num_bits: u8, reg: A64Opnd, disp: i32) -> Self {
+        A64Opnd::Mem(A64Mem::new(num_bits, reg, disp))
     }
 
     /// Convenience function to check if this operand is a register.
@@ -87,23 +92,32 @@ pub const X0_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 0 };
 pub const X1_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 1 };
 pub const X2_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 2 };
 pub const X3_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 3 };
+pub const X4_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 4 };
+pub const X5_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 5 };
 
+pub const X9_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 9 };
+pub const X10_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 10 };
+pub const X11_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 11 };
 pub const X12_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 12 };
 pub const X13_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 13 };
 
+pub const X24_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 24 };
+pub const X25_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 25 };
+pub const X26_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 26 };
+
 // 64-bit registers
 pub const X0: A64Opnd = A64Opnd::Reg(X0_REG);
 pub const X1: A64Opnd = A64Opnd::Reg(X1_REG);
 pub const X2: A64Opnd = A64Opnd::Reg(X2_REG);
 pub const X3: A64Opnd = A64Opnd::Reg(X3_REG);
-pub const X4: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 4 });
-pub const X5: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 5 });
+pub const X4: A64Opnd = A64Opnd::Reg(X4_REG);
+pub const X5: A64Opnd = A64Opnd::Reg(X5_REG);
 pub const X6: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 6 });
 pub const X7: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 7 });
 pub const X8: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 8 });
-pub const X9: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 9 });
-pub const X10: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 10 });
-pub const X11: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 11 });
+pub const X9: A64Opnd = A64Opnd::Reg(X9_REG);
+pub const X10: A64Opnd = A64Opnd::Reg(X10_REG);
+pub const X11: A64Opnd = A64Opnd::Reg(X11_REG);
 pub const X12: A64Opnd = A64Opnd::Reg(X12_REG);
 pub const X13: A64Opnd = A64Opnd::Reg(X13_REG);
 pub const X14: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 14 });
@@ -116,13 +130,14 @@ pub const X20: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 20 });
 pub const X21: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 21 });
 pub const X22: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 22 });
 pub const X23: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 23 });
-pub const X24: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 24 });
-pub const X25: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 25 });
-pub const X26: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 26 });
+pub const X24: A64Opnd = A64Opnd::Reg(X24_REG);
+pub const X25: A64Opnd = A64Opnd::Reg(X25_REG);
+pub const X26: A64Opnd = A64Opnd::Reg(X26_REG);
 pub const X27: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 27 });
 pub const X28: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 28 });
 pub const X29: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 29 });
 pub const X30: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 30 });
+pub const X31: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 31 });
 
 // 32-bit registers
 pub const W0: A64Reg = A64Reg { num_bits: 32, reg_no: 0 };
@@ -156,6 +171,7 @@ pub const W27: A64Reg = A64Reg { num_bits: 32, reg_no: 27 };
 pub const W28: A64Reg = A64Reg { num_bits: 32, reg_no: 28 };
 pub const W29: A64Reg = A64Reg { num_bits: 32, reg_no: 29 };
 pub const W30: A64Reg = A64Reg { num_bits: 32, reg_no: 30 };
+pub const W31: A64Reg = A64Reg { num_bits: 32, reg_no: 31 };
 
 // C argument registers
 pub const C_ARG_REGS: [A64Opnd; 4] = [X0, X1, X2, X3];