diff --git a/crates/js/src/jit/assembler.rs b/crates/js/src/jit/assembler.rs new file mode 100644 index 0000000..3b25123 --- /dev/null +++ b/crates/js/src/jit/assembler.rs @@ -0,0 +1,1500 @@ +//! AArch64 machine code assembler. +//! +//! Emits fixed-width 32-bit AArch64 instructions into a byte buffer. +//! Supports label-based forward references with backpatching. + +use super::buffer::{CodePtr, JitBuffer}; +use super::memory::MemoryError; + +// ── AArch64 register definitions ───────────────────────────────────────────── + +/// A 64-bit general-purpose register (X0–X30, SP/ZR). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Reg(u8); + +impl Reg { + pub const fn new(n: u8) -> Self { + assert!(n < 32, "register index must be 0-31"); + Reg(n) + } + + pub const fn index(self) -> u8 { + self.0 + } +} + +// Arguments / return value +pub const X0: Reg = Reg(0); +pub const X1: Reg = Reg(1); +pub const X2: Reg = Reg(2); +pub const X3: Reg = Reg(3); +pub const X4: Reg = Reg(4); +pub const X5: Reg = Reg(5); +pub const X6: Reg = Reg(6); +pub const X7: Reg = Reg(7); + +// Indirect result location +pub const X8: Reg = Reg(8); + +// Temporaries (caller-saved) +pub const X9: Reg = Reg(9); +pub const X10: Reg = Reg(10); +pub const X11: Reg = Reg(11); +pub const X12: Reg = Reg(12); +pub const X13: Reg = Reg(13); +pub const X14: Reg = Reg(14); +pub const X15: Reg = Reg(15); + +// Intra-procedure-call temporary +pub const X16: Reg = Reg(16); +pub const X17: Reg = Reg(17); + +// Platform register (reserved) +pub const X18: Reg = Reg(18); + +// Callee-saved registers +pub const X19: Reg = Reg(19); +pub const X20: Reg = Reg(20); +pub const X21: Reg = Reg(21); +pub const X22: Reg = Reg(22); +pub const X23: Reg = Reg(23); +pub const X24: Reg = Reg(24); +pub const X25: Reg = Reg(25); +pub const X26: Reg = Reg(26); +pub const X27: Reg = Reg(27); +pub const X28: Reg = Reg(28); + +// Frame pointer +pub const FP: Reg = Reg(29); +// Link register +pub const LR: Reg = Reg(30); +// Stack pointer / zero register (context-dependent) +pub const SP: Reg = Reg(31); +pub const ZR: Reg = Reg(31); + +/// A 64-bit floating-point / SIMD register (D0–D31). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct FpReg(u8); + +impl FpReg { + pub const fn new(n: u8) -> Self { + assert!(n < 32, "FP register index must be 0-31"); + FpReg(n) + } + + pub const fn index(self) -> u8 { + self.0 + } +} + +pub const D0: FpReg = FpReg(0); +pub const D1: FpReg = FpReg(1); +pub const D2: FpReg = FpReg(2); +pub const D3: FpReg = FpReg(3); +pub const D4: FpReg = FpReg(4); +pub const D5: FpReg = FpReg(5); +pub const D6: FpReg = FpReg(6); +pub const D7: FpReg = FpReg(7); + +// ── Condition codes ────────────────────────────────────────────────────────── + +/// AArch64 condition codes for conditional branches. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +pub enum Cond { + EQ = 0b0000, // Equal (Z=1) + NE = 0b0001, // Not equal (Z=0) + CS = 0b0010, // Carry set / unsigned higher or same (C=1) + CC = 0b0011, // Carry clear / unsigned lower (C=0) + MI = 0b0100, // Minus / negative (N=1) + PL = 0b0101, // Plus / positive or zero (N=0) + VS = 0b0110, // Overflow (V=1) + VC = 0b0111, // No overflow (V=0) + HI = 0b1000, // Unsigned higher (C=1 && Z=0) + LS = 0b1001, // Unsigned lower or same (C=0 || Z=1) + GE = 0b1010, // Signed greater or equal (N==V) + LT = 0b1011, // Signed less than (N!=V) + GT = 0b1100, // Signed greater than (Z==0 && N==V) + LE = 0b1101, // Signed less or equal (Z==1 || N!=V) + AL = 0b1110, // Always +} + +// ── Labels ─────────────────────────────────────────────────────────────────── + +/// A label for forward/backward branch references. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Label(u32); + +/// Tracks pending fixups for forward references. +struct PendingFixup { + /// Byte offset in the code buffer where the branch instruction lives. + instruction_offset: usize, + /// The label being referenced. + label: Label, + /// What kind of fixup is needed. + kind: FixupKind, +} + +#[derive(Debug, Clone, Copy)] +enum FixupKind { + /// B/BL: 26-bit signed offset (imm26), in words + Branch26, + /// B.cond/CBZ/CBNZ: 19-bit signed offset (imm19), in words + Branch19, +} + +// ── Assembler ──────────────────────────────────────────────────────────────── + +/// AArch64 machine code assembler. +/// +/// Emits instructions into an internal byte buffer, then finalizes into +/// executable code via a JitBuffer. +pub struct Assembler { + /// Emitted machine code bytes. + code: Vec, + /// Label ID → byte offset (None if not yet bound). + labels: Vec>, + /// Forward references to patch when labels are bound. + fixups: Vec, +} + +impl Default for Assembler { + fn default() -> Self { + Self::new() + } +} + +impl Assembler { + pub fn new() -> Self { + Assembler { + code: Vec::with_capacity(1024), + labels: Vec::new(), + fixups: Vec::new(), + } + } + + /// Current byte offset in the code buffer. + pub fn offset(&self) -> usize { + self.code.len() + } + + /// Get the emitted code as a byte slice. + pub fn code(&self) -> &[u8] { + &self.code + } + + /// Create a new unbound label. + pub fn new_label(&mut self) -> Label { + let id = self.labels.len() as u32; + self.labels.push(None); + Label(id) + } + + /// Bind a label to the current code offset. + /// Patches all pending forward references to this label. + pub fn bind_label(&mut self, label: Label) { + let offset = self.code.len(); + assert!( + self.labels[label.0 as usize].is_none(), + "label already bound" + ); + self.labels[label.0 as usize] = Some(offset); + + // Backpatch pending fixups + let fixups: Vec<_> = self + .fixups + .iter() + .filter(|f| f.label == label) + .map(|f| (f.instruction_offset, f.kind)) + .collect(); + + for (inst_offset, kind) in fixups { + let word_offset = (offset as i64 - inst_offset as i64) / 4; + self.patch_branch(inst_offset, word_offset, kind); + } + self.fixups.retain(|f| f.label != label); + } + + /// Finalize: resolve all labels, emit code into a JitBuffer, return CodePtr. + pub fn finalize(&self, buffer: &mut JitBuffer) -> Result { + assert!(self.fixups.is_empty(), "unresolved forward references"); + buffer.emit_code(&self.code) + } + + /// Finalize and return the raw code bytes (for testing without executable memory). + pub fn finalize_to_vec(&self) -> Vec { + assert!(self.fixups.is_empty(), "unresolved forward references"); + self.code.clone() + } + + // ── Instruction emission helpers ───────────────────────────────────── + + fn emit_u32(&mut self, inst: u32) { + self.code.extend_from_slice(&inst.to_le_bytes()); + } + + fn patch_branch(&mut self, inst_offset: usize, word_offset: i64, kind: FixupKind) { + let mut inst = + u32::from_le_bytes(self.code[inst_offset..inst_offset + 4].try_into().unwrap()); + + match kind { + FixupKind::Branch26 => { + assert!( + (-(1i64 << 25)..(1i64 << 25)).contains(&word_offset), + "branch26 offset out of range: {}", + word_offset + ); + let imm26 = (word_offset as u32) & 0x03FF_FFFF; + inst = (inst & 0xFC00_0000) | imm26; + } + FixupKind::Branch19 => { + assert!( + (-(1i64 << 18)..(1i64 << 18)).contains(&word_offset), + "branch19 offset out of range: {}", + word_offset + ); + let imm19 = (word_offset as u32) & 0x7FFFF; + inst = (inst & 0xFF00_001F) | (imm19 << 5); + } + } + + self.code[inst_offset..inst_offset + 4].copy_from_slice(&inst.to_le_bytes()); + } + + fn emit_branch_to_label(&mut self, label: Label, kind: FixupKind, base_inst: u32) { + if let Some(target) = self.labels[label.0 as usize] { + // Backward reference — patch immediately + let word_offset = (target as i64 - self.code.len() as i64) / 4; + let inst = match kind { + FixupKind::Branch26 => { + let imm26 = (word_offset as u32) & 0x03FF_FFFF; + (base_inst & 0xFC00_0000) | imm26 + } + FixupKind::Branch19 => { + let imm19 = (word_offset as u32) & 0x7FFFF; + (base_inst & 0xFF00_001F) | (imm19 << 5) + } + }; + self.emit_u32(inst); + } else { + // Forward reference — emit placeholder and record fixup + self.fixups.push(PendingFixup { + instruction_offset: self.code.len(), + label, + kind, + }); + self.emit_u32(base_inst); + } + } + + // ── Arithmetic instructions ────────────────────────────────────────── + + /// ADD Xd, Xn, Xm + pub fn add_reg(&mut self, rd: Reg, rn: Reg, rm: Reg) { + // 1_00_01011_00_0_Rm_000000_Rn_Rd + let inst = 0x8B000000 + | ((rm.index() as u32) << 16) + | ((rn.index() as u32) << 5) + | (rd.index() as u32); + self.emit_u32(inst); + } + + /// ADD Xd, Xn, #imm12 + pub fn add_imm(&mut self, rd: Reg, rn: Reg, imm12: u16) { + assert!(imm12 < 4096, "immediate must be < 4096"); + // 1_00_100010_0_imm12_Rn_Rd + let inst = + 0x91000000 | ((imm12 as u32) << 10) | ((rn.index() as u32) << 5) | (rd.index() as u32); + self.emit_u32(inst); + } + + /// SUB Xd, Xn, Xm + pub fn sub_reg(&mut self, rd: Reg, rn: Reg, rm: Reg) { + // 1_10_01011_00_0_Rm_000000_Rn_Rd + let inst = 0xCB000000 + | ((rm.index() as u32) << 16) + | ((rn.index() as u32) << 5) + | (rd.index() as u32); + self.emit_u32(inst); + } + + /// SUB Xd, Xn, #imm12 + pub fn sub_imm(&mut self, rd: Reg, rn: Reg, imm12: u16) { + assert!(imm12 < 4096, "immediate must be < 4096"); + // 1_10_100010_0_imm12_Rn_Rd + let inst = + 0xD1000000 | ((imm12 as u32) << 10) | ((rn.index() as u32) << 5) | (rd.index() as u32); + self.emit_u32(inst); + } + + /// MUL Xd, Xn, Xm (alias for MADD Xd, Xn, Xm, XZR) + pub fn mul(&mut self, rd: Reg, rn: Reg, rm: Reg) { + // 1_00_11011_000_Rm_0_11111_Rn_Rd (MADD with Ra=XZR) + let inst = 0x9B007C00 + | ((rm.index() as u32) << 16) + | ((rn.index() as u32) << 5) + | (rd.index() as u32); + self.emit_u32(inst); + } + + /// SDIV Xd, Xn, Xm + pub fn sdiv(&mut self, rd: Reg, rn: Reg, rm: Reg) { + // 1_00_11010110_Rm_00001_1_Rn_Rd + let inst = 0x9AC00C00 + | ((rm.index() as u32) << 16) + | ((rn.index() as u32) << 5) + | (rd.index() as u32); + self.emit_u32(inst); + } + + // ── Logic instructions ─────────────────────────────────────────────── + + /// AND Xd, Xn, Xm + pub fn and_reg(&mut self, rd: Reg, rn: Reg, rm: Reg) { + // 1_00_01010_00_0_Rm_000000_Rn_Rd + let inst = 0x8A000000 + | ((rm.index() as u32) << 16) + | ((rn.index() as u32) << 5) + | (rd.index() as u32); + self.emit_u32(inst); + } + + /// ORR Xd, Xn, Xm + pub fn orr_reg(&mut self, rd: Reg, rn: Reg, rm: Reg) { + // 1_01_01010_00_0_Rm_000000_Rn_Rd + let inst = 0xAA000000 + | ((rm.index() as u32) << 16) + | ((rn.index() as u32) << 5) + | (rd.index() as u32); + self.emit_u32(inst); + } + + /// EOR Xd, Xn, Xm + pub fn eor_reg(&mut self, rd: Reg, rn: Reg, rm: Reg) { + // 1_10_01010_00_0_Rm_000000_Rn_Rd + let inst = 0xCA000000 + | ((rm.index() as u32) << 16) + | ((rn.index() as u32) << 5) + | (rd.index() as u32); + self.emit_u32(inst); + } + + /// LSL Xd, Xn, Xm (alias for LSLV) + pub fn lsl_reg(&mut self, rd: Reg, rn: Reg, rm: Reg) { + // 1_00_11010110_Rm_00100_0_Rn_Rd + let inst = 0x9AC02000 + | ((rm.index() as u32) << 16) + | ((rn.index() as u32) << 5) + | (rd.index() as u32); + self.emit_u32(inst); + } + + /// LSR Xd, Xn, Xm (alias for LSRV) + pub fn lsr_reg(&mut self, rd: Reg, rn: Reg, rm: Reg) { + // 1_00_11010110_Rm_00100_1_Rn_Rd + let inst = 0x9AC02400 + | ((rm.index() as u32) << 16) + | ((rn.index() as u32) << 5) + | (rd.index() as u32); + self.emit_u32(inst); + } + + /// ASR Xd, Xn, Xm (alias for ASRV) + pub fn asr_reg(&mut self, rd: Reg, rn: Reg, rm: Reg) { + // 1_00_11010110_Rm_00101_0_Rn_Rd + let inst = 0x9AC02800 + | ((rm.index() as u32) << 16) + | ((rn.index() as u32) << 5) + | (rd.index() as u32); + self.emit_u32(inst); + } + + // ── Memory instructions ────────────────────────────────────────────── + + /// LDR Xd, [Xn, Xm] (register offset) + pub fn ldr_reg(&mut self, rt: Reg, rn: Reg, rm: Reg) { + // 11_111_0_00_01_1_Rm_011_0_10_Rn_Rt (LSL #0) + let inst = 0xF8606800 + | ((rm.index() as u32) << 16) + | ((rn.index() as u32) << 5) + | (rt.index() as u32); + self.emit_u32(inst); + } + + /// LDR Xd, [Xn, #imm12] (unsigned offset, scaled by 8) + pub fn ldr_imm(&mut self, rt: Reg, rn: Reg, imm12: u16) { + assert!(imm12.is_multiple_of(8), "offset must be 8-byte aligned"); + let scaled = imm12 / 8; + assert!(scaled < 4096, "scaled offset must be < 4096"); + // 11_111_0_01_01_imm12_Rn_Rt + let inst = + 0xF9400000 | ((scaled as u32) << 10) | ((rn.index() as u32) << 5) | (rt.index() as u32); + self.emit_u32(inst); + } + + /// LDR Xd, [Xn, #simm9]! (pre-index) + pub fn ldr_pre(&mut self, rt: Reg, rn: Reg, simm9: i16) { + assert!( + (-256..=255).contains(&simm9), + "pre-index offset must be -256..255" + ); + let imm9 = (simm9 as u32) & 0x1FF; + // 11_111_0_00_01_0_imm9_11_Rn_Rt + let inst = 0xF8400C00 | (imm9 << 12) | ((rn.index() as u32) << 5) | (rt.index() as u32); + self.emit_u32(inst); + } + + /// LDR Xd, [Xn], #simm9 (post-index) + pub fn ldr_post(&mut self, rt: Reg, rn: Reg, simm9: i16) { + assert!( + (-256..=255).contains(&simm9), + "post-index offset must be -256..255" + ); + let imm9 = (simm9 as u32) & 0x1FF; + // 11_111_0_00_01_0_imm9_01_Rn_Rt + let inst = 0xF8400400 | (imm9 << 12) | ((rn.index() as u32) << 5) | (rt.index() as u32); + self.emit_u32(inst); + } + + /// STR Xd, [Xn, Xm] (register offset) + pub fn str_reg(&mut self, rt: Reg, rn: Reg, rm: Reg) { + // 11_111_0_00_00_1_Rm_011_0_10_Rn_Rt (LSL #0) + let inst = 0xF8206800 + | ((rm.index() as u32) << 16) + | ((rn.index() as u32) << 5) + | (rt.index() as u32); + self.emit_u32(inst); + } + + /// STR Xd, [Xn, #imm12] (unsigned offset, scaled by 8) + pub fn str_imm(&mut self, rt: Reg, rn: Reg, imm12: u16) { + assert!(imm12.is_multiple_of(8), "offset must be 8-byte aligned"); + let scaled = imm12 / 8; + assert!(scaled < 4096, "scaled offset must be < 4096"); + // 11_111_0_01_00_imm12_Rn_Rt + let inst = + 0xF9000000 | ((scaled as u32) << 10) | ((rn.index() as u32) << 5) | (rt.index() as u32); + self.emit_u32(inst); + } + + /// STR Xd, [Xn, #simm9]! (pre-index) + pub fn str_pre(&mut self, rt: Reg, rn: Reg, simm9: i16) { + assert!( + (-256..=255).contains(&simm9), + "pre-index offset must be -256..255" + ); + let imm9 = (simm9 as u32) & 0x1FF; + // 11_111_0_00_00_0_imm9_11_Rn_Rt + let inst = 0xF8000C00 | (imm9 << 12) | ((rn.index() as u32) << 5) | (rt.index() as u32); + self.emit_u32(inst); + } + + /// STR Xd, [Xn], #simm9 (post-index) + pub fn str_post(&mut self, rt: Reg, rn: Reg, simm9: i16) { + assert!( + (-256..=255).contains(&simm9), + "post-index offset must be -256..255" + ); + let imm9 = (simm9 as u32) & 0x1FF; + // 11_111_0_00_00_0_imm9_01_Rn_Rt + let inst = 0xF8000400 | (imm9 << 12) | ((rn.index() as u32) << 5) | (rt.index() as u32); + self.emit_u32(inst); + } + + /// STP Xt1, Xt2, [Xn, #imm7]! (pre-index, store pair) + pub fn stp_pre(&mut self, rt1: Reg, rt2: Reg, rn: Reg, simm7: i16) { + assert!(simm7 % 8 == 0, "offset must be 8-byte aligned"); + let scaled = simm7 / 8; + assert!( + (-64..=63).contains(&scaled), + "scaled offset must be -64..63" + ); + let imm7 = (scaled as u32) & 0x7F; + // 10_101_0_011_imm7_Rt2_Rn_Rt1 + let inst = 0xA9800000 + | (imm7 << 15) + | ((rt2.index() as u32) << 10) + | ((rn.index() as u32) << 5) + | (rt1.index() as u32); + self.emit_u32(inst); + } + + /// LDP Xt1, Xt2, [Xn], #imm7 (post-index, load pair) + pub fn ldp_post(&mut self, rt1: Reg, rt2: Reg, rn: Reg, simm7: i16) { + assert!(simm7 % 8 == 0, "offset must be 8-byte aligned"); + let scaled = simm7 / 8; + assert!( + (-64..=63).contains(&scaled), + "scaled offset must be -64..63" + ); + let imm7 = (scaled as u32) & 0x7F; + // 10_101_0_001_1_imm7_Rt2_Rn_Rt1 + let inst = 0xA8C00000 + | (imm7 << 15) + | ((rt2.index() as u32) << 10) + | ((rn.index() as u32) << 5) + | (rt1.index() as u32); + self.emit_u32(inst); + } + + // ── Branch instructions ────────────────────────────────────────────── + + /// B