jit_assembler/riscv64/
builder.rs

1/// Instruction builder interface for RISC-V 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 RISC-V instructions
11pub struct Riscv64InstructionBuilder {
12    instructions: Vec<Instruction>,
13    #[cfg(feature = "register-tracking")]
14    register_usage: crate::common::register_usage::RegisterUsageInfo<Register>,
15}
16
17impl Riscv64InstructionBuilder {
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    // Register tracking wrapper functions for encode_* functions
63    
64    /// R-type instruction with register tracking: rd = f(rs1, rs2)
65    fn encode_r_type_tracked(&mut self, opcode: u8, rd: Register, funct3: u8, rs1: Register, rs2: Register, funct7: u8) -> Instruction {
66        self.track_written_register(rd);
67        self.track_read_registers(&[rs1, rs2]);
68        encode_r_type(opcode, rd, funct3, rs1, rs2, funct7)
69    }
70    
71    /// I-type instruction with register tracking: rd = f(rs1, imm)
72    fn encode_i_type_tracked(&mut self, opcode: u8, rd: Register, funct3: u8, rs1: Register, imm: i16) -> Instruction {
73        self.track_written_register(rd);
74        self.track_read_register(rs1);
75        encode_i_type(opcode, rd, funct3, rs1, imm)
76    }
77    
78    /// S-type instruction with register tracking: MEM[rs1 + imm] = rs2
79    fn encode_s_type_tracked(&mut self, opcode: u8, funct3: u8, rs1: Register, rs2: Register, imm: i16) -> Instruction {
80        self.track_read_registers(&[rs1, rs2]);
81        encode_s_type(opcode, funct3, rs1, rs2, imm)
82    }
83    
84    /// B-type instruction with register tracking: branch if f(rs1, rs2)
85    fn encode_b_type_tracked(&mut self, opcode: u8, funct3: u8, rs1: Register, rs2: Register, imm: i16) -> Instruction {
86        self.track_read_registers(&[rs1, rs2]);
87        encode_b_type(opcode, funct3, rs1, rs2, imm)
88    }
89    
90    /// U-type instruction with register tracking: rd = imm << 12
91    fn encode_u_type_tracked(&mut self, opcode: u8, rd: Register, imm: u32) -> Instruction {
92        self.track_written_register(rd);
93        encode_u_type(opcode, rd, imm)
94    }
95    
96    /// J-type instruction with register tracking: rd = PC + 4, PC += imm
97    fn encode_j_type_tracked(&mut self, opcode: u8, rd: Register, imm: i32) -> Instruction {
98        self.track_written_register(rd);
99        encode_j_type(opcode, rd, imm)
100    }
101    
102    /// CSR-type instruction with register tracking: rd = CSR, CSR = f(CSR, rs1)
103    fn encode_csr_type_tracked(&mut self, opcode: u8, rd: Register, funct3: u8, rs1: Register, csr: Csr) -> Instruction {
104        self.track_written_register(rd);
105        self.track_read_register(rs1);
106        encode_csr_type(opcode, rd, funct3, rs1, csr)
107    }
108    
109    /// CSR immediate-type instruction with register tracking: rd = CSR, CSR = f(CSR, uimm)
110    fn encode_csr_imm_type_tracked(&mut self, opcode: u8, rd: Register, funct3: u8, uimm: u8, csr: Csr) -> Instruction {
111        self.track_written_register(rd);
112        encode_csr_imm_type(opcode, rd, funct3, uimm, csr)
113    }
114
115    /// Returns a slice of the raw instructions.
116    ///
117    /// This method exposes the internal instruction buffer directly as a slice.
118    /// Prefer using the `instructions()` method from the `InstructionBuilder` trait
119    /// for most use cases, as it provides a higher-level abstraction and is part of
120    /// the public API. Use `raw_instructions` only if you specifically need access
121    /// to the underlying slice for migration or performance reasons.
122    pub fn raw_instructions(&self) -> &[Instruction] {
123        &self.instructions
124    }
125
126    pub fn push(&mut self, instr: Instruction) -> &mut Self {
127        self.instructions.push(instr);
128        self
129    }
130
131    pub fn clear(&mut self) -> &mut Self {
132        self.instructions.clear();
133        #[cfg(feature = "register-tracking")]
134        self.register_usage.clear();
135        self
136    }
137}
138
139impl InstructionBuilder<Instruction> for Riscv64InstructionBuilder {
140    type Register = Register;
141    
142    fn new() -> Self {
143        Self {
144            instructions: Vec::new(),
145            #[cfg(feature = "register-tracking")]
146            register_usage: crate::common::register_usage::RegisterUsageInfo::new(),
147        }
148    }
149
150    fn instructions(&self) -> crate::common::InstructionCollection<Instruction> {
151        crate::common::InstructionCollection::from_slice(&self.instructions)
152    }
153
154    fn push(&mut self, instr: Instruction) {
155        self.instructions.push(instr);
156    }
157
158    fn clear(&mut self) {
159        self.instructions.clear();
160        #[cfg(feature = "register-tracking")]
161        self.register_usage.clear();
162    }
163    
164    #[cfg(feature = "register-tracking")]
165    fn register_usage(&self) -> &crate::common::register_usage::RegisterUsageInfo<Self::Register> {
166        &self.register_usage
167    }
168    
169    #[cfg(feature = "register-tracking")]
170    fn register_usage_mut(&mut self) -> &mut crate::common::register_usage::RegisterUsageInfo<Self::Register> {
171        &mut self.register_usage
172    }
173    
174    /// Create a JIT-compiled function from the assembled instructions (std-only)
175    /// 
176    /// This method converts the assembled instructions into executable machine code
177    /// that can be called directly as a function. The generic type parameter `F`
178    /// specifies the function signature.
179    /// 
180    /// # Safety
181    /// 
182    /// This function is unsafe because:
183    /// - It allocates executable memory
184    /// - It assumes the assembled code follows the correct ABI
185    /// - The caller must ensure the function signature matches the actual code
186    /// 
187    /// # Examples
188    /// 
189    /// ```rust,no_run
190    /// use jit_assembler::riscv64::{reg, Riscv64InstructionBuilder};
191    /// use jit_assembler::common::InstructionBuilder;
192    /// 
193    /// let add_func = unsafe {
194    ///     Riscv64InstructionBuilder::new()
195    ///         .add(reg::A0, reg::A0, reg::A1) // Add first two arguments
196    ///         .ret()
197    ///         .function::<fn(u64, u64) -> u64>()
198    /// }.expect("Failed to create JIT function");
199    /// 
200    /// // Call the JIT function directly (only works on RISC-V hosts)
201    /// let result = add_func.call(10, 20); // Should return 30
202    /// ```
203    #[cfg(feature = "std")]
204    unsafe fn function<F>(&self) -> Result<crate::common::jit::CallableJitFunction<F>, crate::common::jit::JitError> {
205        // Convert instructions to bytes using the new Instructions struct
206        let code = self.instructions().to_bytes();
207        crate::common::jit::CallableJitFunction::<F>::new(&code)
208    }
209
210    #[cfg(feature = "std")]
211    unsafe fn raw_function(&self) -> Result<crate::common::jit::RawCallableJitFunction, crate::common::jit::JitError> {
212        let code = self.instructions().to_bytes();
213        crate::common::jit::RawCallableJitFunction::new(&code)
214    }
215}
216
217impl Riscv64InstructionBuilder {
218    /// Generate CSR read-write instruction
219    pub fn csrrw(&mut self, rd: Register, csr: Csr, rs1: Register) -> &mut Self {
220        let instr = self.encode_csr_type_tracked(opcodes::SYSTEM, rd, system_funct3::CSRRW, rs1, csr);
221        self.push(instr);
222        self
223    }
224
225    /// Generate CSR read-set instruction
226    pub fn csrrs(&mut self, rd: Register, csr: Csr, rs1: Register) -> &mut Self {
227        let instr = self.encode_csr_type_tracked(opcodes::SYSTEM, rd, system_funct3::CSRRS, rs1, csr);
228        self.push(instr);
229        self
230    }
231
232    /// Generate CSR read-clear instruction
233    pub fn csrrc(&mut self, rd: Register, csr: Csr, rs1: Register) -> &mut Self {
234        let instr = self.encode_csr_type_tracked(opcodes::SYSTEM, rd, system_funct3::CSRRC, rs1, csr);
235        self.push(instr);
236        self
237    }
238
239    /// Generate CSR read-write immediate instruction
240    pub fn csrrwi(&mut self, rd: Register, csr: Csr, uimm: u8) -> &mut Self {
241        let instr = self.encode_csr_imm_type_tracked(opcodes::SYSTEM, rd, system_funct3::CSRRWI, uimm, csr);
242        self.push(instr);
243        self
244    }
245
246    /// Generate CSR read-set immediate instruction
247    pub fn csrrsi(&mut self, rd: Register, csr: Csr, uimm: u8) -> &mut Self {
248        let instr = self.encode_csr_imm_type_tracked(opcodes::SYSTEM, rd, system_funct3::CSRRSI, uimm, csr);
249        self.push(instr);
250        self
251    }
252
253    /// Generate CSR read-clear immediate instruction
254    pub fn csrrci(&mut self, rd: Register, csr: Csr, uimm: u8) -> &mut Self {
255        let instr = self.encode_csr_imm_type_tracked(opcodes::SYSTEM, rd, system_funct3::CSRRCI, uimm, csr);
256        self.push(instr);
257        self
258    }
259
260    // Pseudo-instructions for convenience
261
262    /// CSR read (alias for csrrs with rs1=x0)
263    /// This is a common alias in RISC-V assembly
264    pub fn csrr(&mut self, rd: Register, csr: Csr) -> &mut Self {
265        self.csrrs(rd, csr, super::reg::X0)
266    }
267
268    /// CSR write (alias for csrrw with rd=x0)
269    /// This is a common alias in RISC-V assembly for writing to CSR without reading old value
270    pub fn csrw(&mut self, csr: Csr, rs1: Register) -> &mut Self {
271        self.csrrw(super::reg::X0, csr, rs1)
272    }
273
274    /// CSR set (alias for csrrs with rd=x0)
275    /// This is a common alias in RISC-V assembly for setting bits in CSR without reading old value
276    pub fn csrs(&mut self, csr: Csr, rs1: Register) -> &mut Self {
277        self.csrrs(super::reg::X0, csr, rs1)
278    }
279
280    /// CSR clear (alias for csrrc with rd=x0)
281    /// This is a common alias in RISC-V assembly for clearing bits in CSR without reading old value
282    pub fn csrc(&mut self, csr: Csr, rs1: Register) -> &mut Self {
283        self.csrrc(super::reg::X0, csr, rs1)
284    }
285
286    /// CSR write immediate (alias for csrrwi with rd=x0)
287    /// This is a common alias in RISC-V assembly for writing immediate to CSR without reading old value
288    pub fn csrwi(&mut self, csr: Csr, uimm: u8) -> &mut Self {
289        self.csrrwi(super::reg::X0, csr, uimm)
290    }
291
292    /// CSR set immediate (alias for csrrsi with rd=x0)
293    /// This is a common alias in RISC-V assembly for setting bits in CSR with immediate without reading old value
294    pub fn csrsi(&mut self, csr: Csr, uimm: u8) -> &mut Self {
295        self.csrrsi(super::reg::X0, csr, uimm)
296    }
297
298    /// CSR clear immediate (alias for csrrci with rd=x0)
299    /// This is a common alias in RISC-V assembly for clearing bits in CSR with immediate without reading old value
300    pub fn csrci(&mut self, csr: Csr, uimm: u8) -> &mut Self {
301        self.csrrci(super::reg::X0, csr, uimm)
302    }
303
304    // ALU instructions
305
306    /// Generate add instruction
307    pub fn add(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
308        let instr = self.encode_r_type_tracked(opcodes::OP, rd, alu_funct3::ADD_SUB, rs1, rs2, 0x0);
309        self.push(instr);
310        self
311    }
312    /// Generate add immediate instruction
313    pub fn addi(&mut self, rd: Register, rs1: Register, imm: i16) -> &mut Self {
314        let instr = self.encode_i_type_tracked(opcodes::OP_IMM, rd, alu_funct3::ADD_SUB, rs1, imm);
315        self.push(instr);
316        self
317    }
318
319
320    /// Generate subtract instruction
321    pub fn sub(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
322        let instr = self.encode_r_type_tracked(opcodes::OP, rd, alu_funct3::ADD_SUB, rs1, rs2, 0x20);
323        self.push(instr);
324        self
325    }
326
327
328    /// Generate subtract immediate instruction
329    pub fn subi(&mut self, rd: Register, rs1: Register, imm: i16) -> &mut Self {
330        let instr = self.encode_i_type_tracked(opcodes::OP_IMM, rd, alu_funct3::ADD_SUB, rs1, -imm);
331        self.push(instr);
332        self
333    }
334
335    /// Generate XOR instruction
336    pub fn xor(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
337        let instr = self.encode_r_type_tracked(opcodes::OP, rd, alu_funct3::XOR, rs1, rs2, 0x0);
338        self.push(instr);
339        self
340    }
341
342    /// Generate XOR immediate instruction
343    pub fn xori(&mut self, rd: Register, rs1: Register, imm: i16) -> &mut Self {
344        let instr = self.encode_i_type_tracked(opcodes::OP_IMM, rd, alu_funct3::XOR, rs1, imm);
345        self.push(instr);
346        self
347    }
348
349    /// Generate OR instruction
350    pub fn or(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
351        let instr = self.encode_r_type_tracked(opcodes::OP, rd, alu_funct3::OR, rs1, rs2, 0x0);
352        self.push(instr);
353        self
354    }
355
356    /// Generate OR immediate instruction
357    pub fn ori(&mut self, rd: Register, rs1: Register, imm: i16) -> &mut Self {
358        let instr = self.encode_i_type_tracked(opcodes::OP_IMM, rd, alu_funct3::OR, rs1, imm);
359        self.push(instr);
360        self
361    }
362
363    /// Generate Set Less Than instruction
364    pub fn slt(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
365        let instr = self.encode_r_type_tracked(opcodes::OP, rd, alu_funct3::SLT, rs1, rs2, 0x0);
366        self.push(instr);
367        self
368    }
369
370    /// Generate Set Less Than immediate instruction
371    pub fn slti(&mut self, rd: Register, rs1: Register, imm: i16) -> &mut Self {
372        let instr = self.encode_i_type_tracked(opcodes::OP_IMM, rd, alu_funct3::SLT, rs1, imm);
373        self.push(instr);
374        self
375    }
376
377    /// Generate Set Less Than Unsigned instruction
378    pub fn sltu(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
379        let instr = self.encode_r_type_tracked(opcodes::OP, rd, alu_funct3::SLTU, rs1, rs2, 0x0);
380        self.push(instr);
381        self
382    }
383
384    /// Generate Set Less Than immediate Unsigned instruction
385    pub fn sltiu(&mut self, rd: Register, rs1: Register, imm: i16) -> &mut Self {
386        let instr = self.encode_i_type_tracked(opcodes::OP_IMM, rd, alu_funct3::SLTU, rs1, imm);
387        self.push(instr);
388        self
389    }
390
391    /// Generate SLL (Shift Left Logical) instruction
392    pub fn sll(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
393        let instr = self.encode_r_type_tracked(opcodes::OP, rd, alu_funct3::SLL, rs1, rs2, 0x0);
394        self.push(instr);
395        self
396    }
397
398    /// Generate SRL (Shift Right Logical) instruction
399    pub fn srl(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
400        let instr = self.encode_r_type_tracked(opcodes::OP, rd, alu_funct3::SRL_SRA, rs1, rs2, 0x0);
401        self.push(instr);
402        self
403    }
404
405    /// Generate SRA (Shift Right Arithmetic) instruction
406    pub fn sra(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
407        let instr = self.encode_r_type_tracked(opcodes::OP, rd, alu_funct3::SRL_SRA, rs1, rs2, 0x20);
408        self.push(instr);
409        self
410    }
411
412    /// Generate SLLI (Shift Left Logical Immediate) instruction
413    pub fn slli(&mut self, rd: Register, rs1: Register, shamt: u8) -> &mut Self {
414        let instr = self.encode_i_type_tracked(opcodes::OP_IMM, rd, alu_funct3::SLL, rs1, shamt as i16);
415        self.push(instr);
416        self
417    }
418
419    /// Generate SRLI (Shift Right Logical Immediate) instruction
420    pub fn srli(&mut self, rd: Register, rs1: Register, shamt: u8) -> &mut Self {
421        let instr = self.encode_i_type_tracked(opcodes::OP_IMM, rd, alu_funct3::SRL_SRA, rs1, shamt as i16);
422        self.push(instr);
423        self
424    }
425
426    /// Generate SRAI (Shift Right Arithmetic Immediate) instruction
427    pub fn srai(&mut self, rd: Register, rs1: Register, shamt: u8) -> &mut Self {
428        let instr = self.encode_i_type_tracked(opcodes::OP_IMM, rd, alu_funct3::SRL_SRA, rs1, (shamt as i16) | 0x400);
429        self.push(instr);
430        self
431    }
432
433    /// Generate AND instruction
434    pub fn and(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
435        let instr = self.encode_r_type_tracked(opcodes::OP, rd, alu_funct3::AND, rs1, rs2, 0x0);
436        self.push(instr);
437        self
438    }
439
440    /// Generate AND immediate instruction
441    pub fn andi(&mut self, rd: Register, rs1: Register, imm: i16) -> &mut Self {
442        let instr = self.encode_i_type_tracked(opcodes::OP_IMM, rd, alu_funct3::AND, rs1, imm);
443        self.push(instr);
444        self
445    }
446
447    // M Extension (Multiply/Divide) instructions
448
449    /// Generate MUL (Multiply) instruction
450    /// Performs signed multiplication and returns the lower 64 bits of the result
451    pub fn mul(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
452        let instr = self.encode_r_type_tracked(opcodes::OP, rd, m_funct3::MUL, rs1, rs2, m_funct7::M_EXT);
453        self.push(instr);
454        self
455    }
456
457    /// Generate MULH (Multiply High) instruction
458    /// Performs signed × signed multiplication and returns the upper 64 bits of the result
459    pub fn mulh(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
460        let instr = self.encode_r_type_tracked(opcodes::OP, rd, m_funct3::MULH, rs1, rs2, m_funct7::M_EXT);
461        self.push(instr);
462        self
463    }
464
465    /// Generate MULHSU (Multiply High Signed × Unsigned) instruction
466    /// Performs signed × unsigned multiplication and returns the upper 64 bits of the result
467    pub fn mulhsu(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
468        let instr = self.encode_r_type_tracked(opcodes::OP, rd, m_funct3::MULHSU, rs1, rs2, m_funct7::M_EXT);
469        self.push(instr);
470        self
471    }
472
473    /// Generate MULHU (Multiply High Unsigned) instruction
474    /// Performs unsigned × unsigned multiplication and returns the upper 64 bits of the result
475    pub fn mulhu(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
476        let instr = self.encode_r_type_tracked(opcodes::OP, rd, m_funct3::MULHU, rs1, rs2, m_funct7::M_EXT);
477        self.push(instr);
478        self
479    }
480
481    /// Generate DIV (Divide) instruction
482    /// Performs signed division: rs1 ÷ rs2
483    pub fn div(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
484        let instr = self.encode_r_type_tracked(opcodes::OP, rd, m_funct3::DIV, rs1, rs2, m_funct7::M_EXT);
485        self.push(instr);
486        self
487    }
488
489    /// Generate DIVU (Divide Unsigned) instruction
490    /// Performs unsigned division: rs1 ÷ rs2
491    pub fn divu(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
492        let instr = self.encode_r_type_tracked(opcodes::OP, rd, m_funct3::DIVU, rs1, rs2, m_funct7::M_EXT);
493        self.push(instr);
494        self
495    }
496
497    /// Generate REM (Remainder) instruction
498    /// Computes signed remainder: rs1 % rs2
499    pub fn rem(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
500        let instr = self.encode_r_type_tracked(opcodes::OP, rd, m_funct3::REM, rs1, rs2, m_funct7::M_EXT);
501        self.push(instr);
502        self
503    }
504
505    /// Generate REMU (Remainder Unsigned) instruction
506    /// Computes unsigned remainder: rs1 % rs2
507    pub fn remu(&mut self, rd: Register, rs1: Register, rs2: Register) -> &mut Self {
508        let instr = self.encode_r_type_tracked(opcodes::OP, rd, m_funct3::REMU, rs1, rs2, m_funct7::M_EXT);
509        self.push(instr);
510        self
511    }
512
513    // Load/Store instructions
514
515    /// Generate LD (Load Doubleword) instruction
516    pub fn ld(&mut self, rd: Register, rs1: Register, imm: i16) -> &mut Self {
517        let instr = self.encode_i_type_tracked(opcodes::LOAD, rd, load_funct3::LD, rs1, imm);
518        self.push(instr);
519        self
520    }
521
522    /// Generate LW (Load Word) instruction
523    pub fn lw(&mut self, rd: Register, rs1: Register, imm: i16) -> &mut Self {
524        let instr = self.encode_i_type_tracked(opcodes::LOAD, rd, load_funct3::LW, rs1, imm);
525        self.push(instr);
526        self
527    }
528
529    /// Generate LH (Load Halfword) instruction
530    pub fn lh(&mut self, rd: Register, rs1: Register, imm: i16) -> &mut Self {
531        let instr = self.encode_i_type_tracked(opcodes::LOAD, rd, load_funct3::LH, rs1, imm);
532        self.push(instr);
533        self
534    }
535
536    /// Generate LB (Load Byte) instruction
537    pub fn lb(&mut self, rd: Register, rs1: Register, imm: i16) -> &mut Self {
538        let instr = self.encode_i_type_tracked(opcodes::LOAD, rd, load_funct3::LB, rs1, imm);
539        self.push(instr);
540        self
541    }
542
543    /// Generate LBU (Load Byte Unsigned) instruction
544    pub fn lbu(&mut self, rd: Register, rs1: Register, imm: i16) -> &mut Self {
545        let instr = self.encode_i_type_tracked(opcodes::LOAD, rd, load_funct3::LBU, rs1, imm);
546        self.push(instr);
547        self
548    }
549
550    /// Generate LHU (Load Halfword Unsigned) instruction
551    pub fn lhu(&mut self, rd: Register, rs1: Register, imm: i16) -> &mut Self {
552        let instr = self.encode_i_type_tracked(opcodes::LOAD, rd, load_funct3::LHU, rs1, imm);
553        self.push(instr);
554        self
555    }
556
557    /// Generate LWU (Load Word Unsigned) instruction
558    pub fn lwu(&mut self, rd: Register, rs1: Register, imm: i16) -> &mut Self {
559        let instr = self.encode_i_type_tracked(opcodes::LOAD, rd, load_funct3::LWU, rs1, imm);
560        self.push(instr);
561        self
562    }
563
564    /// Load immediate value into register (handles large immediates)
565    /// This is a common pseudo-instruction in RISC-V assembly
566    pub fn li(&mut self, rd: Register, imm: i32) -> &mut Self {
567        // LUI loads the upper 20 bits, ADDI adds the lower 12 bits
568        let upper = (imm + 0x800) >> 12; // Round up if lower 12 bits are negative
569        let lower = imm & 0xfff;
570        if upper != 0 {
571            self.lui(rd, upper as u32);
572        }
573        if lower != 0 || upper == 0 {
574            // When upper == 0, rd has not been initialized by lui,
575            // so we must use x0 (zero register) as source to properly initialize rd
576            if upper == 0 {
577                self.addi(rd, super::reg::X0, lower as i16);
578            } else {
579                self.addi(rd, rd, lower as i16);
580            }
581        }
582        self
583    }
584
585    /// Generate SD (Store Doubleword) instruction
586    pub fn sd(&mut self, rs1: Register, rs2: Register, imm: i16) -> &mut Self {
587        let instr = self.encode_s_type_tracked(opcodes::STORE, store_funct3::SD, rs1, rs2, imm);
588        self.push(instr);
589        self
590    }
591
592    /// Generate SW (Store Word) instruction
593    pub fn sw(&mut self, rs1: Register, rs2: Register, imm: i16) -> &mut Self {
594        let instr = self.encode_s_type_tracked(opcodes::STORE, store_funct3::SW, rs1, rs2, imm);
595        self.push(instr);
596        self
597    }
598
599    /// Generate SH (Store Halfword) instruction
600    pub fn sh(&mut self, rs1: Register, rs2: Register, imm: i16) -> &mut Self {
601        let instr = self.encode_s_type_tracked(opcodes::STORE, store_funct3::SH, rs1, rs2, imm);
602        self.push(instr);
603        self
604    }
605
606    /// Generate SB (Store Byte) instruction
607    pub fn sb(&mut self, rs1: Register, rs2: Register, imm: i16) -> &mut Self {
608        let instr = self.encode_s_type_tracked(opcodes::STORE, store_funct3::SB, rs1, rs2, imm);
609        self.push(instr);
610        self
611    }
612
613    /// Generate LUI (Load Upper Immediate) instruction
614    pub fn lui(&mut self, rd: Register, imm: u32) -> &mut Self {
615        let instr = self.encode_u_type_tracked(opcodes::LUI, rd, imm);
616        self.push(instr);
617        self
618    }
619
620    /// Generate AUIPC (Add Upper Immediate to PC) instruction
621    pub fn auipc(&mut self, rd: Register, imm: u32) -> &mut Self {
622        let instr = self.encode_u_type_tracked(opcodes::AUIPC, rd, imm);
623        self.push(instr);
624        self
625    }
626
627    // Control flow instructions
628
629    /// Generate JAL (Jump and Link) instruction
630    pub fn jal(&mut self, rd: Register, imm: i32) -> &mut Self {
631        let instr = self.encode_j_type_tracked(opcodes::JAL, rd, imm);
632        self.push(instr);
633        self
634    }
635
636    /// Generate JALR (Jump and Link Register) instruction
637    pub fn jalr(&mut self, rd: Register, rs1: Register, imm: i16) -> &mut Self {
638        let instr = self.encode_i_type_tracked(opcodes::JALR, rd, 0x0, rs1, imm);
639        self.push(instr);
640        self
641    }
642
643    /// Generate BEQ (Branch if Equal) instruction
644    pub fn beq(&mut self, rs1: Register, rs2: Register, imm: i16) -> &mut Self {
645        let instr = self.encode_b_type_tracked(opcodes::BRANCH, branch_funct3::BEQ, rs1, rs2, imm);
646        self.push(instr);
647        self
648    }
649
650    /// Generate BNE (Branch if Not Equal) instruction
651    pub fn bne(&mut self, rs1: Register, rs2: Register, imm: i16) -> &mut Self {
652        let instr = self.encode_b_type_tracked(opcodes::BRANCH, branch_funct3::BNE, rs1, rs2, imm);
653        self.push(instr);
654        self
655    }
656
657    /// Generate BLT (Branch if Less Than) instruction
658    pub fn blt(&mut self, rs1: Register, rs2: Register, imm: i16) -> &mut Self {
659        let instr = self.encode_b_type_tracked(opcodes::BRANCH, branch_funct3::BLT, rs1, rs2, imm);
660        self.push(instr);
661        self
662    }
663
664    /// Generate BGE (Branch if Greater or Equal) instruction
665    pub fn bge(&mut self, rs1: Register, rs2: Register, imm: i16) -> &mut Self {
666        let instr = self.encode_b_type_tracked(opcodes::BRANCH, branch_funct3::BGE, rs1, rs2, imm);
667        self.push(instr);
668        self
669    }
670
671    /// Generate BLTU (Branch if Less Than Unsigned) instruction
672    pub fn bltu(&mut self, rs1: Register, rs2: Register, imm: i16) -> &mut Self {
673        let instr = self.encode_b_type_tracked(opcodes::BRANCH, branch_funct3::BLTU, rs1, rs2, imm);
674        self.push(instr);
675        self
676    }
677
678    /// Generate BGEU (Branch if Greater or Equal Unsigned) instruction
679    pub fn bgeu(&mut self, rs1: Register, rs2: Register, imm: i16) -> &mut Self {
680        let instr = self.encode_b_type_tracked(opcodes::BRANCH, branch_funct3::BGEU, rs1, rs2, imm);
681        self.push(instr);
682        self
683    }
684
685    // Privileged instructions
686
687    /// Supervisor return instruction
688    /// Returns from supervisor mode to user mode or previous privilege level
689    /// This is a privileged instruction that can only be executed in supervisor mode or higher
690    pub fn sret(&mut self) -> &mut Self {
691        let instr = super::instruction::encode_privileged_type(
692            super::instruction::opcodes::SYSTEM,
693            super::instruction::privileged_funct12::SRET
694        );
695        self.push(instr);
696        self
697    }
698
699    /// Machine return instruction
700    /// Returns from machine mode to previous privilege level
701    /// This is a privileged instruction that can only be executed in machine mode
702    pub fn mret(&mut self) -> &mut Self {
703        let instr = super::instruction::encode_privileged_type(
704            super::instruction::opcodes::SYSTEM,
705            super::instruction::privileged_funct12::MRET
706        );
707        self.push(instr);
708        self
709    }
710
711    /// Environment call instruction
712    /// Generates a system call to the execution environment
713    pub fn ecall(&mut self) -> &mut Self {
714        let instr = super::instruction::encode_privileged_type(
715            super::instruction::opcodes::SYSTEM,
716            super::instruction::privileged_funct12::ECALL
717        );
718        self.push(instr);
719        self
720    }
721
722    /// Environment break instruction
723    /// Generates a breakpoint exception
724    pub fn ebreak(&mut self) -> &mut Self {
725        let instr = super::instruction::encode_privileged_type(
726            super::instruction::opcodes::SYSTEM,
727            super::instruction::privileged_funct12::EBREAK
728        );
729        self.push(instr);
730        self
731    }
732
733    /// Wait for interrupt instruction
734    /// Puts the processor in a low-power state until an interrupt occurs
735    /// This is a privileged instruction
736    pub fn wfi(&mut self) -> &mut Self {
737        let instr = super::instruction::encode_privileged_type(
738            super::instruction::opcodes::SYSTEM,
739            super::instruction::privileged_funct12::WFI
740        );
741        self.push(instr);
742        self
743    }
744
745    /// Return instruction (alias for jalr x0, x1, 0)
746    /// This is a common alias in RISC-V assembly for returning from a function
747    pub fn ret(&mut self) -> &mut Self {
748        self.jalr(super::reg::X0, super::reg::X1, 0)
749    }
750}