jit_assembler/aarch64/
instruction.rs

1/// AArch64 instruction formats and encoding
2use core::fmt;
3use crate::common::{
4    Instruction as InstructionTrait,
5    Register as RegisterTrait,
6};
7
8#[cfg(feature = "std")]
9use std::vec::Vec;
10#[cfg(not(feature = "std"))]
11use alloc::vec::Vec;
12
13/// AArch64 register representation (32 general-purpose registers)
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
15pub struct Register(pub u8);
16
17impl Register {
18    pub const fn new(reg: u8) -> Self {
19        Self(reg)
20    }
21
22    pub fn value(self) -> u8 {
23        self.0
24    }
25}
26
27impl RegisterTrait for Register {
28    fn id(&self) -> u32 {
29        self.0 as u32
30    }
31    
32    fn abi_class(&self) -> crate::common::AbiClass {
33        use crate::common::AbiClass;
34        
35        match self.0 {
36            // Caller-saved registers (do not need to be preserved across calls)
37            0..=7 => AbiClass::CallerSaved,     // X0-X7: Argument/return value registers
38            8..=15 => AbiClass::CallerSaved,    // X8-X15: Caller-saved temporary registers
39            16..=17 => AbiClass::CallerSaved,   // X16-X17: Intra-procedure-call registers
40            18 => AbiClass::CallerSaved,        // X18: Platform register (caller-saved on most platforms)
41            
42            // Callee-saved registers (must be preserved across calls)
43            19..=28 => AbiClass::CalleeSaved,   // X19-X28: Callee-saved registers
44            
45            // Special-purpose registers
46            29 => AbiClass::Special,  // X29: Frame pointer (FP)
47            30 => AbiClass::Special,  // X30: Link register (LR)
48            31 => AbiClass::Special,  // X31: Stack pointer (SP) or zero register (XZR)
49
50            // Default to Special for any unhandled registers
51            _ => AbiClass::Special,
52        }
53    }
54}
55
56/// AArch64 instruction representation (32-bit fixed-width instructions)
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub struct Instruction(pub u32);
59
60impl Instruction {
61    /// Create a new 32-bit instruction
62    pub fn new(value: u32) -> Self {
63        Self(value)
64    }
65
66    /// Get the instruction value as u32
67    pub fn value(self) -> u32 {
68        self.0
69    }
70
71    /// Get the instruction as bytes (little-endian)
72    pub fn bytes(self) -> Vec<u8> {
73        self.0.to_le_bytes().to_vec()
74    }
75}
76
77impl InstructionTrait for Instruction {
78    fn value(&self) -> u64 {
79        self.0 as u64
80    }
81    
82    fn bytes(&self) -> Vec<u8> {
83        self.0.to_le_bytes().to_vec()
84    }
85    
86    fn size(&self) -> usize {
87        4
88    }
89}
90
91impl fmt::Display for Instruction {
92    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
93        write!(f, "0x{:08x}", self.0)
94    }
95}
96
97/// Data Processing - Register instruction encoding (3-operand)  
98pub fn encode_add_sub_reg(sf: u8, op: u8, s: u8, rm: Register, imm6: u8, rn: Register, rd: Register) -> Instruction {
99    // ADD/SUB (shifted register) encoding according to AArch64 ISA
100    // 31: sf (0=32-bit, 1=64-bit)
101    // 30: op (0=ADD, 1=SUB) 
102    // 29: s (0=don't set flags, 1=set flags)
103    // 28-24: 01011 (fixed for shifted register)
104    // 23-22: shift (00=LSL, 01=LSR, 10=ASR, 11=reserved)
105    // 21: 0 (fixed)
106    // 20-16: Rm 
107    // 15-10: imm6 (shift amount)
108    // 9-5: Rn
109    // 4-0: Rd
110    let instr = ((sf as u32) << 31) |
111                ((op as u32) << 30) |
112                ((s as u32) << 29) |
113                (0b01011 << 24) |     // Fixed bits for shifted register format
114                (0b00 << 22) |        // shift = 00 (LSL)
115                (0 << 21) |           // Fixed bit
116                ((rm.value() as u32) << 16) |
117                ((imm6 as u32) << 10) |
118                ((rn.value() as u32) << 5) |
119                (rd.value() as u32);
120    Instruction::new(instr)
121}
122
123/// Data Processing - Immediate instruction encoding (ADD/SUB immediate)
124pub fn encode_add_sub_imm(sf: u8, op: u8, s: u8, sh: u8, imm12: u16, rn: Register, rd: Register) -> Instruction {
125    let instr = ((sf as u32) << 31) |
126                ((op as u32) << 30) |
127                ((s as u32) << 29) |
128                (0b10001 << 24) |  // Fixed bits for data processing immediate
129                ((sh as u32) << 22) |
130                ((imm12 as u32) << 10) |
131                ((rn.value() as u32) << 5) |
132                (rd.value() as u32);
133    Instruction::new(instr)
134}
135
136/// Logical instruction encoding (register)
137pub fn encode_logical_reg(sf: u8, opc: u8, shift: u8, n: u8, rm: Register, imm6: u8, rn: Register, rd: Register) -> Instruction {
138    let instr = ((sf as u32) << 31) |
139                ((opc as u32) << 29) |
140                (0b01010 << 24) |  // Fixed bits for logical register
141                ((shift as u32) << 22) |
142                ((n as u32) << 21) |
143                ((rm.value() as u32) << 16) |
144                ((imm6 as u32) << 10) |
145                ((rn.value() as u32) << 5) |
146                (rd.value() as u32);
147    Instruction::new(instr)
148}
149
150/// Multiply instruction encoding
151pub fn encode_multiply(sf: u8, op31: u8, rm: Register, o0: u8, ra: Register, rn: Register, rd: Register) -> Instruction {
152    // For MUL x0, x1, x2 -> encoding should be 0x9b027c20
153    // AArch64 Data Processing -- 3 source format:
154    // sf | op54 | 11011 | op31 | Rm | o0 | Ra | Rn | Rd
155    // 1  | 00   | 11011 | 000  | Rm | 0  | Ra | Rn | Rd
156    let instr = ((sf as u32) << 31) |         // sf (size flag)
157                (0b00 << 29) |                // op54 = 00 
158                (0b11011 << 24) |             // Fixed: 11011
159                ((op31 as u32) << 21) |       // op31 (operation)
160                ((rm.value() as u32) << 16) | // Rm (source register 2)
161                ((o0 as u32) << 15) |         // o0 (operation variant)
162                ((ra.value() as u32) << 10) | // Ra (accumulator, XZR for MUL)
163                ((rn.value() as u32) << 5) |  // Rn (source register 1)
164                (rd.value() as u32);          // Rd (destination)
165    Instruction::new(instr)
166}
167
168/// Division instruction encoding - Data Processing (2 source)
169pub fn encode_divide(sf: u8, opcode: u8, rm: Register, _o0: u8, rn: Register, rd: Register) -> Instruction {
170    // Data-processing (2 source) format:
171    // sf | 0 | S | 11010110 | Rm | opcode | Rn | Rd
172    // For UDIV: opcode = 000010
173    // For SDIV: opcode = 000011
174    let instr = ((sf as u32) << 31) |         // sf (64-bit)
175                (0b0 << 30) |                 // op bit
176                (0b0 << 29) |                 // S bit  
177                (0b11010110 << 21) |          // Fixed bits for data processing 2-source
178                ((rm.value() as u32) << 16) | // Rm register
179                ((opcode as u32) << 10) |     // opcode (UDIV=000010, SDIV=000011)
180                ((rn.value() as u32) << 5) |  // Rn register
181                (rd.value() as u32);          // Rd register
182    Instruction::new(instr)
183}
184
185/// Move instruction encoding (ORR with XZR)
186pub fn encode_move_reg(sf: u8, rm: Register, rd: Register) -> Instruction {
187    // MOV Rd, Rm -> ORR Rd, XZR, Rm
188    encode_logical_reg(sf, 0b01, 0b00, 0, rm, 0, Register::new(31), rd)
189}
190
191/// Return instruction encoding (RET)
192pub fn encode_ret(rn: Register) -> Instruction {
193    // RET instruction -> encoding should be 0xd65f03c0 for ret (X30 implied)
194    // Based on GNU assembler output analysis
195    // Pattern: 11010110 01011111 00000011 11000000 = 0xd65f03c0
196    // The 0x5f pattern is bits 16-22: 0101111 (bit pattern: 0x5f at positions 16-22)
197    let instr = 0xd6000000 |           // Base RET instruction pattern
198                (0x5f << 16) |         // bits 22-16: 0101111 (matches GNU assembler exactly)
199                ((rn.value() as u32) << 5);  // bits 9-5: Rn
200    Instruction::new(instr)
201}
202
203/// Branch register instruction encoding (BR)
204pub fn encode_branch_reg(opc: u8, op2: u8, op3: u8, rn: Register, op4: u8) -> Instruction {
205    let instr = (0b1101011 << 25) |
206                ((opc as u32) << 21) |
207                ((op2 as u32) << 16) |
208                ((op3 as u32) << 10) |
209                ((rn.value() as u32) << 5) |
210                (op4 as u32);
211    Instruction::new(instr)
212}
213
214/// Common registers
215pub mod reg {
216    use super::Register;
217    
218    // Standard register names (X0-X30)
219    pub const X0: Register = Register::new(0);
220    pub const X1: Register = Register::new(1);
221    pub const X2: Register = Register::new(2);
222    pub const X3: Register = Register::new(3);
223    pub const X4: Register = Register::new(4);
224    pub const X5: Register = Register::new(5);
225    pub const X6: Register = Register::new(6);
226    pub const X7: Register = Register::new(7);
227    pub const X8: Register = Register::new(8);
228    pub const X9: Register = Register::new(9);
229    pub const X10: Register = Register::new(10);
230    pub const X11: Register = Register::new(11);
231    pub const X12: Register = Register::new(12);
232    pub const X13: Register = Register::new(13);
233    pub const X14: Register = Register::new(14);
234    pub const X15: Register = Register::new(15);
235    pub const X16: Register = Register::new(16);
236    pub const X17: Register = Register::new(17);
237    pub const X18: Register = Register::new(18);
238    pub const X19: Register = Register::new(19);
239    pub const X20: Register = Register::new(20);
240    pub const X21: Register = Register::new(21);
241    pub const X22: Register = Register::new(22);
242    pub const X23: Register = Register::new(23);
243    pub const X24: Register = Register::new(24);
244    pub const X25: Register = Register::new(25);
245    pub const X26: Register = Register::new(26);
246    pub const X27: Register = Register::new(27);
247    pub const X28: Register = Register::new(28);
248    pub const X29: Register = Register::new(29);
249    pub const X30: Register = Register::new(30);
250    
251    // X31 is special - it's SP in some contexts, XZR/WZR in others
252    pub const SP: Register = Register::new(31);   // Stack pointer
253    pub const XZR: Register = Register::new(31);  // Zero register (64-bit)
254    pub const WZR: Register = Register::new(31);  // Zero register (32-bit)
255
256    // AArch64 ABI register aliases
257    pub const FP: Register = X29;    // Frame pointer
258    pub const LR: Register = X30;    // Link register
259}
260
261/// MOVZ instruction encoding - Move immediate with zero
262pub fn encode_movz(sf: u8, hw: u8, imm16: u16, rd: Register) -> Instruction {
263    // MOVZ Xd, #imm, LSL #(hw*16)
264    // sf | 10 | 100101 | hw | imm16 | Rd
265    // 1  | 10 | 100101 | hw | imm16 | Rd (for 64-bit)
266    let instr = ((sf as u32) << 31) |         // sf (64-bit)
267                (0b10 << 29) |                // opc = 10 for MOVZ
268                (0b100101 << 23) |            // Fixed bits for move wide immediate
269                ((hw as u32) << 21) |         // hw (shift amount / 16)
270                ((imm16 as u32) << 5) |       // imm16 (16-bit immediate)
271                (rd.value() as u32);          // Rd (destination register)
272    Instruction::new(instr)
273}
274
275/// MOVK instruction encoding - Move immediate with keep
276pub fn encode_movk(sf: u8, hw: u8, imm16: u16, rd: Register) -> Instruction {
277    // MOVK Xd, #imm, LSL #(hw*16)
278    // sf | 11 | 100101 | hw | imm16 | Rd
279    // 1  | 11 | 100101 | hw | imm16 | Rd (for 64-bit)
280    let instr = ((sf as u32) << 31) |         // sf (64-bit)
281                (0b11 << 29) |                // opc = 11 for MOVK
282                (0b100101 << 23) |            // Fixed bits for move wide immediate
283                ((hw as u32) << 21) |         // hw (shift amount / 16)
284                ((imm16 as u32) << 5) |       // imm16 (16-bit immediate)
285                (rd.value() as u32);          // Rd (destination register)
286    Instruction::new(instr)
287}