jit_assembler/aarch64/
builder.rs

1/// Instruction builder interface for AArch64 assembly generation
2use super::instruction::*;
3use crate::common::InstructionBuilder;
4
5#[cfg(feature = "std")]
6use std::vec::Vec;
7#[cfg(not(feature = "std"))]
8use alloc::vec::Vec;
9
10/// Instruction builder for generating AArch64 instructions
11pub struct Aarch64InstructionBuilder {
12    instructions: Vec<Instruction>,
13    #[cfg(feature = "register-tracking")]
14    register_usage: crate::common::register_usage::RegisterUsageInfo<Register>,
15}
16
17impl Aarch64InstructionBuilder {
18    pub fn new() -> Self {
19        Self {
20            instructions: Vec::new(),
21            #[cfg(feature = "register-tracking")]
22            register_usage: crate::common::register_usage::RegisterUsageInfo::new(),
23        }
24    }
25    
26    /// Track a written register (register-tracking feature only)
27    #[cfg(feature = "register-tracking")]
28    fn track_written_register(&mut self, reg: Register) {
29        self.register_usage.add_written_register(reg);
30    }
31    
32    /// Track a read register (register-tracking feature only)
33    #[cfg(feature = "register-tracking")]
34    fn track_read_register(&mut self, reg: Register) {
35        self.register_usage.add_read_register(reg);
36    }
37    
38    /// Track multiple read registers at once (register-tracking feature only)
39    #[cfg(feature = "register-tracking")]
40    fn track_read_registers(&mut self, regs: &[Register]) {
41        for &reg in regs {
42            self.register_usage.add_read_register(reg);
43        }
44    }
45    
46    /// No-op versions for when register-tracking is disabled
47    #[cfg(not(feature = "register-tracking"))]
48    fn track_written_register(&mut self, _reg: Register) {
49        // No-op
50    }
51    
52    #[cfg(not(feature = "register-tracking"))]
53    fn track_read_register(&mut self, _reg: Register) {
54        // No-op
55    }
56    
57    #[cfg(not(feature = "register-tracking"))]
58    fn track_read_registers(&mut self, _regs: &[Register]) {
59        // No-op
60    }
61
62    /// Returns a slice of the raw instructions.
63    ///
64    /// This method exposes the internal instruction buffer directly as a slice.
65    /// Prefer using the `instructions()` method from the `InstructionBuilder` trait
66    /// for most use cases, as it provides a higher-level abstraction and is part of
67    /// the public API. Use `raw_instructions` only if you specifically need access
68    /// to the underlying slice for migration or performance reasons.
69    pub fn raw_instructions(&self) -> &[Instruction] {
70        &self.instructions
71    }
72
73    pub fn push(&mut self, instr: Instruction) -> &mut Self {
74        self.instructions.push(instr);
75        self
76    }
77
78    pub fn clear(&mut self) -> &mut Self {
79        self.instructions.clear();
80        #[cfg(feature = "register-tracking")]
81        self.register_usage.clear();
82        self
83    }
84
85    // Arithmetic instructions
86
87    /// Generate ADD instruction (64-bit register)
88    /// ADD Xd, Xn, Xm
89    pub fn add(&mut self, rd: Register, rn: Register, rm: Register) -> &mut Self {
90        self.track_written_register(rd);
91        self.track_read_registers(&[rn, rm]);
92        let instr = encode_add_sub_reg(1, 0, 0, rm, 0, rn, rd); // sf=1 (64-bit), op=0 (ADD), s=0 (no flags)
93        self.push(instr);
94        self
95    }
96
97    /// Generate ADD immediate instruction (64-bit)
98    /// ADD Xd, Xn, #imm
99    pub fn addi(&mut self, rd: Register, rn: Register, imm: u16) -> &mut Self {
100        self.track_written_register(rd);
101        self.track_read_register(rn);
102        let instr = encode_add_sub_imm(1, 0, 0, 0, imm, rn, rd); // sf=1, op=0 (ADD), s=0, sh=0
103        self.push(instr);
104        self
105    }
106
107    /// Generate SUB instruction (64-bit register)
108    /// SUB Xd, Xn, Xm
109    pub fn sub(&mut self, rd: Register, rn: Register, rm: Register) -> &mut Self {
110        self.track_written_register(rd);
111        self.track_read_registers(&[rn, rm]);
112        let instr = encode_add_sub_reg(1, 1, 0, rm, 0, rn, rd); // sf=1, op=1 (SUB), s=0
113        self.push(instr);
114        self
115    }
116
117    /// Generate SUB immediate instruction (64-bit)
118    /// SUB Xd, Xn, #imm
119    pub fn subi(&mut self, rd: Register, rn: Register, imm: u16) -> &mut Self {
120        self.track_written_register(rd);
121        self.track_read_register(rn);
122        let instr = encode_add_sub_imm(1, 1, 0, 0, imm, rn, rd); // sf=1, op=1 (SUB), s=0, sh=0
123        self.push(instr);
124        self
125    }
126
127    /// Generate MUL instruction (64-bit)
128    /// MUL Xd, Xn, Xm (equivalent to MADD Xd, Xn, Xm, XZR)
129    pub fn mul(&mut self, rd: Register, rn: Register, rm: Register) -> &mut Self {
130        self.track_written_register(rd);
131        self.track_read_registers(&[rn, rm]);
132        let instr = encode_multiply(1, 0b000, rm, 0, Register::new(31), rn, rd); // sf=1, op31=000, o0=0, ra=XZR
133        self.push(instr);
134        self
135    }
136
137    /// Generate MSUB instruction (64-bit multiply-subtract)
138    /// MSUB Xd, Xn, Xm, Xa -> Xd = Xa - (Xn * Xm)
139    pub fn msub(&mut self, rd: Register, rn: Register, rm: Register, ra: Register) -> &mut Self {
140        self.track_written_register(rd);
141        self.track_read_registers(&[rn, rm, ra]);
142        let instr = encode_multiply(1, 0b000, rm, 1, ra, rn, rd); // sf=1, op31=000, o0=1 (MSUB), ra=minuend
143        self.push(instr);
144        self
145    }
146
147    /// Generate UDIV instruction (64-bit unsigned division)
148    /// UDIV Xd, Xn, Xm
149    pub fn udiv(&mut self, rd: Register, rn: Register, rm: Register) -> &mut Self {
150        self.track_written_register(rd);
151        self.track_read_registers(&[rn, rm]);
152        let instr = encode_divide(1, 0b000010, rm, 1, rn, rd); // sf=1, opcode=000010 for UDIV
153        self.push(instr);
154        self
155    }
156
157    /// Generate SDIV instruction (64-bit signed division)
158    /// SDIV Xd, Xn, Xm  
159    pub fn sdiv(&mut self, rd: Register, rn: Register, rm: Register) -> &mut Self {
160        self.track_written_register(rd);
161        self.track_read_registers(&[rn, rm]);
162        let instr = encode_divide(1, 0b000011, rm, 1, rn, rd); // sf=1, opcode=000011 for SDIV
163        self.push(instr);
164        self
165    }
166
167    // DEPRECATED: urem was using a hardcoded temporary register (X17) which is unsafe
168    // Users should implement remainder manually using udiv + msub with their own temp register
169    // Example: 
170    //   .udiv(temp_reg, dividend, divisor)  // temp = dividend / divisor  
171    //   .msub(result, temp, divisor, dividend)  // result = dividend - (temp * divisor)
172    //
173    // /// Generate remainder operation using MSUB after division  
174    // /// This creates a sequence: UDIV tmp, rn, rm; MSUB rd, tmp, rm, rn
175    // /// Result: rd = rn - (rn / rm) * rm = rn % rm
176    // pub fn urem(&mut self, rd: Register, rn: Register, rm: Register) -> &mut Self {
177    //     // We need a temporary register - use X17 (IP1) which is caller-saved
178    //     let tmp = reg::X17;
179    //     
180    //     // UDIV tmp, rn, rm
181    //     self.track_read_registers(&[rn, rm]);
182    //     let div_instr = encode_divide(1, 0b000010, rm, 1, rn, tmp); // opcode=000010 for UDIV
183    //     self.push(div_instr);
184    //     
185    //     // MSUB rd, tmp, rm, rn  (rd = rn - tmp * rm)
186    //     self.track_written_register(rd);
187    //     self.track_read_registers(&[tmp, rm, rn]);
188    //     let msub_instr = encode_multiply(1, 0b000, rm, 1, rn, tmp, rd); // MSUB: op31=000, o0=1
189    //     self.push(msub_instr);
190    //     self
191    // }
192
193    /// Generate ORR instruction (logical OR, 64-bit)
194    /// ORR Xd, Xn, Xm
195    pub fn or(&mut self, rd: Register, rn: Register, rm: Register) -> &mut Self {
196        self.track_written_register(rd);
197        self.track_read_registers(&[rn, rm]);
198        let instr = encode_logical_reg(1, 0b01, 0b00, 0, rm, 0, rn, rd); // sf=1, opc=01 (ORR), shift=00, n=0
199        self.push(instr);
200        self
201    }
202
203    /// Generate AND instruction (logical AND, 64-bit)
204    /// AND Xd, Xn, Xm
205    pub fn and(&mut self, rd: Register, rn: Register, rm: Register) -> &mut Self {
206        self.track_written_register(rd);
207        self.track_read_registers(&[rn, rm]);
208        let instr = encode_logical_reg(1, 0b00, 0b00, 0, rm, 0, rn, rd); // sf=1, opc=00 (AND), shift=00, n=0
209        self.push(instr);
210        self
211    }
212
213    /// Generate EOR instruction (logical XOR, 64-bit)
214    /// EOR Xd, Xn, Xm
215    pub fn xor(&mut self, rd: Register, rn: Register, rm: Register) -> &mut Self {
216        self.track_written_register(rd);
217        self.track_read_registers(&[rn, rm]);
218        let instr = encode_logical_reg(1, 0b10, 0b00, 0, rm, 0, rn, rd); // sf=1, opc=10 (EOR), shift=00, n=0
219        self.push(instr);
220        self
221    }
222
223    /// Generate MOV instruction (move register to register)
224    /// MOV Xd, Xm (implemented as ORR Xd, XZR, Xm)
225    pub fn mov(&mut self, rd: Register, rm: Register) -> &mut Self {
226        self.track_written_register(rd);
227        self.track_read_register(rm);
228        let instr = encode_move_reg(1, rm, rd); // sf=1 (64-bit)
229        self.push(instr);
230        self
231    }
232
233    /// Generate RET instruction (return)
234    /// RET (returns to X30/LR by default)
235    pub fn ret(&mut self) -> &mut Self {
236        self.track_read_register(reg::LR);
237        let instr = encode_ret(reg::LR); // Return to LR (X30)
238        self.push(instr);
239        self
240    }
241
242    /// Generate RET instruction with specific register
243    /// RET Xn
244    pub fn ret_reg(&mut self, rn: Register) -> &mut Self {
245        self.track_read_register(rn);
246        let instr = encode_ret(rn);
247        self.push(instr);
248        self
249    }
250
251    /// Generate MOVZ instruction (Move with Zero)
252    /// MOVZ Xd, #imm16, LSL #(shift*16)
253    pub fn movz(&mut self, rd: Register, imm16: u16, shift: u8) -> &mut Self {
254        self.track_written_register(rd);
255        let instr = encode_movz(1, shift, imm16, rd); // sf=1 for 64-bit
256        self.push(instr);
257        self
258    }
259
260    /// Generate MOVK instruction (Move with Keep)
261    /// MOVK Xd, #imm16, LSL #(shift*16)
262    pub fn movk(&mut self, rd: Register, imm16: u16, shift: u8) -> &mut Self {
263        self.track_written_register(rd);
264        self.track_read_register(rd); // MOVK modifies existing register
265        let instr = encode_movk(1, shift, imm16, rd); // sf=1 for 64-bit
266        self.push(instr);
267        self
268    }
269
270    /// Generate immediate move instruction using MOVZ/MOVK for efficient 64-bit loads
271    pub fn mov_imm(&mut self, rd: Register, imm: u64) -> &mut Self {
272        if imm == 0 {
273            // Special case: move zero register
274            self.mov(rd, reg::XZR);
275            return self;
276        }
277
278        // Break down the 64-bit immediate into 16-bit chunks
279        let chunk0 = (imm & 0xFFFF) as u16;         // Bits 0-15
280        let chunk1 = ((imm >> 16) & 0xFFFF) as u16; // Bits 16-31  
281        let chunk2 = ((imm >> 32) & 0xFFFF) as u16; // Bits 32-47
282        let chunk3 = ((imm >> 48) & 0xFFFF) as u16; // Bits 48-63
283
284        // Find the first non-zero chunk to use MOVZ
285        let mut first_movz_done = false;
286        
287        if chunk0 != 0 {
288            self.movz(rd, chunk0, 0); // LSL #0
289            first_movz_done = true;
290        }
291        
292        if chunk1 != 0 {
293            if first_movz_done {
294                self.movk(rd, chunk1, 1); // LSL #16
295            } else {
296                self.movz(rd, chunk1, 1); // LSL #16
297                first_movz_done = true;
298            }
299        }
300        
301        if chunk2 != 0 {
302            if first_movz_done {
303                self.movk(rd, chunk2, 2); // LSL #32
304            } else {
305                self.movz(rd, chunk2, 2); // LSL #32
306                first_movz_done = true;
307            }
308        }
309        
310        if chunk3 != 0 {
311            if first_movz_done {
312                self.movk(rd, chunk3, 3); // LSL #48
313            } else {
314                self.movz(rd, chunk3, 3); // LSL #48
315                first_movz_done = true;
316            }
317        }
318
319        // If all chunks were zero (shouldn't happen due to early return), move zero
320        if !first_movz_done {
321            self.mov(rd, reg::XZR);
322        }
323
324        self
325    }
326
327    /// Generate left shift instruction (using multiply by power of 2)
328    /// This is a simplified implementation - real AArch64 has LSL instruction
329    pub fn shl(&mut self, rd: Register, rn: Register, shift: u8) -> &mut Self {
330        if shift == 0 {
331            self.mov(rd, rn);
332        } else if shift <= 6 {
333            // Use multiply by 2^shift for small shifts
334            let multiplier = 1u16 << shift;
335            let temp_reg = reg::X17;
336            self.addi(temp_reg, reg::XZR, multiplier);
337            self.mul(rd, rn, temp_reg);
338        } else {
339            // For larger shifts, we'd need proper shift instructions
340            // For now, just copy the register
341            self.mov(rd, rn);
342        }
343        self
344    }
345}
346
347impl InstructionBuilder<Instruction> for Aarch64InstructionBuilder {
348    type Register = Register;
349    
350    fn new() -> Self {
351        Self {
352            instructions: Vec::new(),
353            #[cfg(feature = "register-tracking")]
354            register_usage: crate::common::register_usage::RegisterUsageInfo::new(),
355        }
356    }
357
358    fn instructions(&self) -> crate::common::InstructionCollection<Instruction> {
359        crate::common::InstructionCollection::from_slice(&self.instructions)
360    }
361
362    fn push(&mut self, instr: Instruction) {
363        self.instructions.push(instr);
364    }
365
366    fn clear(&mut self) {
367        self.instructions.clear();
368        #[cfg(feature = "register-tracking")]
369        self.register_usage.clear();
370    }
371    
372    #[cfg(feature = "register-tracking")]
373    fn register_usage(&self) -> &crate::common::register_usage::RegisterUsageInfo<Self::Register> {
374        &self.register_usage
375    }
376    
377    #[cfg(feature = "register-tracking")]
378    fn register_usage_mut(&mut self) -> &mut crate::common::register_usage::RegisterUsageInfo<Self::Register> {
379        &mut self.register_usage
380    }
381    
382    /// Create a JIT-compiled function from the assembled instructions (std-only)
383    /// 
384    /// This method converts the assembled instructions into executable machine code
385    /// that can be called directly as a function. The generic type parameter `F`
386    /// specifies the function signature.
387    /// 
388    /// # Safety
389    /// 
390    /// This function is unsafe because:
391    /// - It allocates executable memory
392    /// - It assumes the assembled code follows the correct ABI
393    /// - The caller must ensure the function signature matches the actual code
394    /// 
395    /// # Examples
396    /// 
397    /// ```rust,no_run
398    /// use jit_assembler::aarch64::{reg, Aarch64InstructionBuilder};
399    /// use jit_assembler::common::InstructionBuilder;
400    /// 
401    /// let add_func = unsafe {
402    ///     Aarch64InstructionBuilder::new()
403    ///         .add(reg::X0, reg::X0, reg::X1) // Add first two arguments
404    ///         .ret()
405    ///         .function::<fn(u64, u64) -> u64>()
406    /// }.expect("Failed to create JIT function");
407    /// 
408    /// // Call the JIT function directly (only works on AArch64 hosts)
409    /// // let result = add_func.call(10, 20); // Returns 30
410    /// ```
411    #[cfg(feature = "std")]
412    unsafe fn function<F>(&self) -> Result<crate::common::jit::CallableJitFunction<F>, crate::common::jit::JitError> {
413        let bytes = self.instructions().to_bytes();
414        crate::common::jit::CallableJitFunction::<F>::new(&bytes)
415    }
416    
417    #[cfg(feature = "std")]
418    unsafe fn raw_function(&self) -> Result<crate::common::jit::RawCallableJitFunction, crate::common::jit::JitError> {
419        let bytes = self.instructions().to_bytes();
420        crate::common::jit::RawCallableJitFunction::new(&bytes)
421    }
422}
jit_assembler/aarch64/builder.rs

jit_assembler/aarch64/
builder.rs
