* Arm64 Beginnings (https://github.com/Shopify/ruby/pull/291)

* Initial setup for aarch64

* ADDS and SUBS

* ADD and SUB for immediates

* Revert moved code

* Documentation

* Rename Arm64* to A64*

* Comments on shift types

* Share sig_imm_size and unsig_imm_size

5 files changed, 485 insertions, 0 deletions

diff --git a/yjit/src/asm/arm64/inst/data_processing_immediate.rs b/yjit/src/asm/arm64/inst/data_processing_immediate.rs
new file mode 100644
index 0000000000..12498848b2
--- /dev/null
+++ b/yjit/src/asm/arm64/inst/data_processing_immediate.rs
@@ -0,0 +1,198 @@
+use super::{
+    super::opnd::*,
+    family::Family,
+    sf::Sf
+};
+
+/// The operation being performed by this instruction.
+enum Op {
+    Add = 0b0,
+    Sub = 0b1
+}
+
+// Whether or not to update the flags when this instruction is performed.
+enum S {
+    LeaveFlags = 0b0,
+    UpdateFlags = 0b1
+}
+
+/// How much to shift the immediate by.
+enum Shift {
+    LSL0 = 0b0, // no shift
+    LSL12 = 0b1 // logical shift left by 12 bits
+}
+
+/// The struct that represents an A64 data processing -- immediate instruction
+/// that can be encoded.
+///
+/// Add/subtract (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  0  1    0                                                                                |
+/// | sf op  S                       sh imm12.................................... rn.............. rd.............. |
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+///
+pub struct DataProcessingImmediate {
+    /// Whether or not this instruction is operating on 64-bit operands.
+    sf: Sf,
+
+    /// The opcode for this instruction.
+    op: Op,
+
+    /// Whether or not to update the flags when this instruction is performed.
+    s: S,
+
+    /// How much to shift the immediate by.
+    shift: Shift,
+
+    /// The value of the immediate.
+    imm12: u16,
+
+    /// The register number of the first operand register.
+    rn: u8,
+
+    /// The register number of the destination register.
+    rd: u8
+}
+
+impl DataProcessingImmediate {
+    /// ADD (immediate)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/ADD--immediate---Add--immediate--?lang=en
+    pub fn add(rd: &A64Opnd, rn: &A64Opnd, imm12: &A64Opnd) -> Self {
+        let (rd, rn, imm12) = Self::unwrap(rd, rn, imm12);
+
+        Self {
+            sf: rd.num_bits.into(),
+            op: Op::Add,
+            s: S::LeaveFlags,
+            shift: Shift::LSL0,
+            imm12: imm12.value as u16,
+            rn: rn.reg_no,
+            rd: rd.reg_no
+        }
+    }
+
+    /// ADDS (immediate, set flags)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/ADDS--immediate---Add--immediate---setting-flags-?lang=en
+    pub fn adds(rd: &A64Opnd, rn: &A64Opnd, imm12: &A64Opnd) -> Self {
+        let (rd, rn, imm12) = Self::unwrap(rd, rn, imm12);
+
+        Self {
+            sf: rd.num_bits.into(),
+            op: Op::Add,
+            s: S::UpdateFlags,
+            shift: Shift::LSL0,
+            imm12: imm12.value as u16,
+            rn: rn.reg_no,
+            rd: rd.reg_no
+        }
+    }
+
+    /// SUB (immediate)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/SUB--immediate---Subtract--immediate--?lang=en
+    pub fn sub(rd: &A64Opnd, rn: &A64Opnd, imm12: &A64Opnd) -> Self {
+        let (rd, rn, imm12) = Self::unwrap(rd, rn, imm12);
+
+        Self {
+            sf: rd.num_bits.into(),
+            op: Op::Sub,
+            s: S::LeaveFlags,
+            shift: Shift::LSL0,
+            imm12: imm12.value as u16,
+            rn: rn.reg_no,
+            rd: rd.reg_no
+        }
+    }
+
+    /// SUBS (immediate, set flags)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/SUBS--immediate---Subtract--immediate---setting-flags-?lang=en
+    pub fn subs(rd: &A64Opnd, rn: &A64Opnd, imm12: &A64Opnd) -> Self {
+        let (rd, rn, imm12) = Self::unwrap(rd, rn, imm12);
+
+        Self {
+            sf: rd.num_bits.into(),
+            op: Op::Sub,
+            s: S::UpdateFlags,
+            shift: Shift::LSL0,
+            imm12: imm12.value as u16,
+            rn: rn.reg_no,
+            rd: rd.reg_no
+        }
+    }
+
+    /// Extract out two registers and an immediate from the given operands.
+    /// Panic if any of the operands do not match the expected type or size.
+    fn unwrap<'a>(rd: &'a A64Opnd, rn: &'a A64Opnd, imm12: &'a A64Opnd) -> (&'a A64Reg, &'a A64Reg, &'a A64UImm) {
+        match (rd, rn, imm12) {
+            (A64Opnd::Reg(rd), A64Opnd::Reg(rn), A64Opnd::UImm(imm12)) => {
+                assert!(rd.num_bits == rn.num_bits, "Both rd and rn operands to a data processing immediate instruction must be of the same size.");
+                assert!(imm12.num_bits <= 12, "The immediate operand to a data processing immediate instruction must be 12 bits or less.");
+                (rd, rn, imm12)
+            },
+            _ => {
+                panic!("Expected 2 register operands and an immediate operand for a data processing immediate instruction.");
+            }
+        }
+    }
+}
+
+impl From<DataProcessingImmediate> for u32 {
+    /// Convert a data processing instruction into a 32-bit value.
+    fn from(inst: DataProcessingImmediate) -> Self {
+        0
+        | (inst.sf as u32).wrapping_shl(31)
+        | (inst.op as u32).wrapping_shl(30)
+        | (inst.s as u32).wrapping_shl(29)
+        | (Family::DataProcessingImmediate as u32).wrapping_shl(25)
+        | (0b1 << 24)
+        | (inst.shift as u32).wrapping_shl(22)
+        | (inst.imm12 as u32).wrapping_shl(10)
+        | (inst.rn as u32).wrapping_shl(5)
+        | inst.rd as u32
+    }
+}
+
+impl From<DataProcessingImmediate> for [u8; 4] {
+    /// Convert a data processing instruction into a 4 byte array.
+    fn from(inst: DataProcessingImmediate) -> [u8; 4] {
+        let result: u32 = inst.into();
+        result.to_le_bytes()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_add() {
+        let uimm12 = A64Opnd::new_uimm(7);
+        let inst = DataProcessingImmediate::add(&X0, &X1, &uimm12);
+        let result: u32 = inst.into();
+        assert_eq!(0x91001c20, result);
+    }
+
+    #[test]
+    fn test_adds() {
+        let uimm12 = A64Opnd::new_uimm(7);
+        let inst = DataProcessingImmediate::adds(&X0, &X1, &uimm12);
+        let result: u32 = inst.into();
+        assert_eq!(0xb1001c20, result);
+    }
+
+    #[test]
+    fn test_sub() {
+        let uimm12 = A64Opnd::new_uimm(7);
+        let inst = DataProcessingImmediate::sub(&X0, &X1, &uimm12);
+        let result: u32 = inst.into();
+        assert_eq!(0xd1001c20, result);
+    }
+
+    #[test]
+    fn test_subs() {
+        let uimm12 = A64Opnd::new_uimm(7);
+        let inst = DataProcessingImmediate::subs(&X0, &X1, &uimm12);
+        let result: u32 = inst.into();
+        assert_eq!(0xf1001c20, result);
+    }
+}
diff --git a/yjit/src/asm/arm64/inst/data_processing_register.rs b/yjit/src/asm/arm64/inst/data_processing_register.rs
new file mode 100644
index 0000000000..6203034e3f
--- /dev/null
+++ b/yjit/src/asm/arm64/inst/data_processing_register.rs
@@ -0,0 +1,202 @@
+use super::{
+    super::opnd::*,
+    family::Family,
+    sf::Sf
+};
+
+/// The operation being performed by this instruction.
+enum Op {
+    Add = 0b0,
+    Sub = 0b1
+}
+
+// Whether or not to update the flags when this instruction is performed.
+enum S {
+    LeaveFlags = 0b0,
+    UpdateFlags = 0b1
+}
+
+/// The type of shift to perform on the second operand register.
+enum Shift {
+    LSL = 0b00, // logical shift left (unsigned)
+    LSR = 0b01, // logical shift right (unsigned)
+    ASR = 0b10  // arithmetic shift right (signed)
+}
+
+/// The struct that represents an A64 data processing -- register instruction
+/// that can be encoded.
+///
+/// Add/subtract (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  1          0                                                                          |
+/// | sf op  S                    shift    rm..............   imm6............... rn.............. rd.............. |
+/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+
+///
+pub struct DataProcessingRegister {
+    /// Whether or not this instruction is operating on 64-bit operands.
+    sf: Sf,
+
+    /// The opcode for this instruction.
+    op: Op,
+
+    /// Whether or not to update the flags when this instruction is performed.
+    s: S,
+
+    /// The type of shift to perform on the second operand register.
+    shift: Shift,
+
+    /// The register number of the second operand register.
+    rm: u8,
+
+    /// The amount to shift the second operand register by.
+    imm6: u8,
+
+    /// The register number of the first operand register.
+    rn: u8,
+
+    /// The register number of the destination register.
+    rd: u8
+}
+
+impl DataProcessingRegister {
+    /// ADD (shifted register)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/ADD--shifted-register---Add--shifted-register--?lang=en
+    pub fn add(rd: &A64Opnd, rn: &A64Opnd, rm: &A64Opnd) -> Self {
+        let (rd, rn, rm) = Self::unwrap(rd, rn, rm);
+
+        Self {
+            sf: rd.num_bits.into(),
+            op: Op::Add,
+            s: S::LeaveFlags,
+            shift: Shift::LSL,
+            rm: rm.reg_no,
+            imm6: 0,
+            rn: rn.reg_no,
+            rd: rd.reg_no
+        }
+    }
+
+    /// ADDS (shifted register, set flags)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/ADDS--shifted-register---Add--shifted-register---setting-flags-?lang=en
+    pub fn adds(rd: &A64Opnd, rn: &A64Opnd, rm: &A64Opnd) -> Self {
+        let (rd, rn, rm) = Self::unwrap(rd, rn, rm);
+
+        Self {
+            sf: rd.num_bits.into(),
+            op: Op::Add,
+            s: S::UpdateFlags,
+            shift: Shift::LSL,
+            rm: rm.reg_no,
+            imm6: 0,
+            rn: rn.reg_no,
+            rd: rd.reg_no
+        }
+    }
+
+    /// SUB (shifted register)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/SUB--shifted-register---Subtract--shifted-register--?lang=en
+    pub fn sub(rd: &A64Opnd, rn: &A64Opnd, rm: &A64Opnd) -> Self {
+        let (rd, rn, rm) = Self::unwrap(rd, rn, rm);
+
+        Self {
+            sf: rd.num_bits.into(),
+            op: Op::Sub,
+            s: S::LeaveFlags,
+            shift: Shift::LSL,
+            rm: rm.reg_no,
+            imm6: 0,
+            rn: rn.reg_no,
+            rd: rd.reg_no
+        }
+    }
+
+    /// SUBS (shifted register, set flags)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/SUBS--shifted-register---Subtract--shifted-register---setting-flags-?lang=en
+    pub fn subs(rd: &A64Opnd, rn: &A64Opnd, rm: &A64Opnd) -> Self {
+        let (rd, rn, rm) = Self::unwrap(rd, rn, rm);
+
+        Self {
+            sf: rd.num_bits.into(),
+            op: Op::Sub,
+            s: S::UpdateFlags,
+            shift: Shift::LSL,
+            rm: rm.reg_no,
+            imm6: 0,
+            rn: rn.reg_no,
+            rd: rd.reg_no
+        }
+    }
+
+    /// Extract out three registers from the given operands. Panic if any of the
+    /// operands are not registers or if they are not the same size.
+    fn unwrap<'a>(rd: &'a A64Opnd, rn: &'a A64Opnd, rm: &'a A64Opnd) -> (&'a A64Reg, &'a A64Reg, &'a A64Reg) {
+        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 to a data processing register instruction must be of the same size.");
+                (rd, rn, rm)
+            },
+            _ => {
+                panic!("Expected 3 register operands for a data processing register instruction.");
+            }
+        }
+    }
+}
+
+impl From<DataProcessingRegister> for u32 {
+    /// Convert a data processing instruction into a 32-bit value.
+    fn from(inst: DataProcessingRegister) -> Self {
+        0
+        | (inst.sf as u32).wrapping_shl(31)
+        | (inst.op as u32).wrapping_shl(30)
+        | (inst.s as u32).wrapping_shl(29)
+        | (Family::DataProcessingRegister as u32).wrapping_shl(25)
+        | (0b1 << 24)
+        | (inst.shift as u32).wrapping_shl(22)
+        | (inst.rm as u32).wrapping_shl(16)
+        | (inst.imm6 as u32).wrapping_shl(10)
+        | (inst.rn as u32).wrapping_shl(5)
+        | inst.rd as u32
+    }
+}
+
+impl From<DataProcessingRegister> for [u8; 4] {
+    /// Convert a data processing instruction into a 4 byte array.
+    fn from(inst: DataProcessingRegister) -> [u8; 4] {
+        let result: u32 = inst.into();
+        result.to_le_bytes()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_add() {
+        let inst = DataProcessingRegister::add(&X0, &X1, &X2);
+        let result: u32 = inst.into();
+        assert_eq!(0x8b020020, result);
+    }
+
+    #[test]
+    fn test_adds() {
+        let inst = DataProcessingRegister::adds(&X0, &X1, &X2);
+        let result: u32 = inst.into();
+        assert_eq!(0xab020020, result);
+    }
+
+    #[test]
+    fn test_sub() {
+        let inst = DataProcessingRegister::sub(&X0, &X1, &X2);
+        let result: u32 = inst.into();
+        assert_eq!(0xcb020020, result);
+    }
+
+    #[test]
+    fn test_subs() {
+        let inst = DataProcessingRegister::subs(&X0, &X1, &X2);
+        let result: u32 = inst.into();
+        assert_eq!(0xeb020020, result);
+    }
+}
diff --git a/yjit/src/asm/arm64/inst/family.rs b/yjit/src/asm/arm64/inst/family.rs
new file mode 100644
index 0000000000..ff5a335406
--- /dev/null
+++ b/yjit/src/asm/arm64/inst/family.rs
@@ -0,0 +1,34 @@
+/// These are the top-level encodings. They're effectively the family of
+/// instructions, as each instruction within those groups shares these same
+/// bits (28-25).
+///
+/// In the documentation, you can see that some of the bits are
+/// optional (e.g., x1x0 for loads and stores). We represent that here as 0100
+/// since we're bitwise ORing the family into the resulting encoding.
+///
+/// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding?lang=en
+pub enum Family {
+    /// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/Reserved?lang=en
+    Reserved = 0b0000,
+
+    /// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/SME-encodings?lang=en
+    SMEEncodings = 0b0001,
+
+    /// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/SVE-encodings?lang=en
+    SVEEncodings = 0b0010,
+
+    /// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/Data-Processing----Immediate?lang=en
+    DataProcessingImmediate = 0b1000,
+
+    /// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/Branches--Exception-Generating-and-System-instructions?lang=en
+    BranchesAndSystem = 0b1010,
+
+    /// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/Loads-and-Stores?lang=en
+    LoadsAndStores = 0b0100,
+
+    /// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/Data-Processing----Register?lang=en
+    DataProcessingRegister = 0b0101,
+
+    /// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/Data-Processing----Scalar-Floating-Point-and-Advanced-SIMD?lang=en
+    DataProcessingScalar = 0b0111
+}
diff --git a/yjit/src/asm/arm64/inst/mod.rs b/yjit/src/asm/arm64/inst/mod.rs
new file mode 100644
index 0000000000..5a0e148ff9
--- /dev/null
+++ b/yjit/src/asm/arm64/inst/mod.rs
@@ -0,0 +1,32 @@
+mod data_processing_immediate;
+mod data_processing_register;
+mod family;
+mod sf;
+
+use data_processing_immediate::DataProcessingImmediate;
+use data_processing_register::DataProcessingRegister;
+
+use crate::asm::CodeBlock;
+use super::opnd::*;
+
+/// ADD
+pub fn add(cb: &mut CodeBlock, rd: &A64Opnd, rn: &A64Opnd, rm: &A64Opnd) {
+    let bytes: [u8; 4] = match rm {
+        A64Opnd::UImm(_) => DataProcessingImmediate::add(rd, rn, rm).into(),
+        A64Opnd::Reg(_) => DataProcessingRegister::add(rd, rn, rm).into(),
+        _ => panic!("Invalid operand combination to add.")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
+/// SUB
+pub fn sub(cb: &mut CodeBlock, rd: &A64Opnd, rn: &A64Opnd, rm: &A64Opnd) {
+    let bytes: [u8; 4] = match rm {
+        A64Opnd::UImm(_) => DataProcessingImmediate::sub(rd, rn, rm).into(),
+        A64Opnd::Reg(_) => DataProcessingRegister::sub(rd, rn, rm).into(),
+        _ => panic!("Invalid operand combination to add.")
+    };
+
+    cb.write_bytes(&bytes);
+}
diff --git a/yjit/src/asm/arm64/inst/sf.rs b/yjit/src/asm/arm64/inst/sf.rs
new file mode 100644
index 0000000000..c2fd33302c
--- /dev/null
+++ b/yjit/src/asm/arm64/inst/sf.rs
@@ -0,0 +1,19 @@
+/// This is commonly the top-most bit in the encoding of the instruction, and
+/// represents whether register operands should be treated as 64-bit registers
+/// or 32-bit registers.
+pub enum Sf {
+    Sf32 = 0b0,
+    Sf64 = 0b1
+}
+
+/// A convenience function so that we can convert the number of bits of an
+/// register operand directly into an Sf enum variant.
+impl From<u8> for Sf {
+    fn from(num_bits: u8) -> Self {
+        match num_bits {
+            64 => Sf::Sf64,
+            32 => Sf::Sf32,
+            _ => panic!("Invalid number of bits: {}", num_bits)
+        }
+    }
+}