# jit-assembler
A multi-architecture JIT assembler library for runtime code generation that works on any host architecture.
## Features
- **Multi-architecture support**: Generate machine code for different target architectures
- **Host-independent**: Runs on any host architecture (x86_64, ARM64, etc.) to generate target code
- **No-std compatible**: Works in both `std` and `no_std` environments
- **Type-safe**: Leverages Rust's type system for safe instruction generation
- **Dual API**: Both macro-based DSL and builder pattern for different use cases
- **IDE-friendly**: Full autocomplete and type checking support
- **JIT execution**: Direct execution of assembled code as functions (std-only)
- **Register usage tracking**: Analyze register usage patterns for optimization and ABI compliance (`register-tracking` feature)
## Supported Architectures
- **RISC-V 64-bit** (`riscv` feature, enabled by default)
- **AArch64** (`aarch64` feature, enabled by default) - Basic arithmetic and logical operations
- **x86-64** (`x86_64` feature) - Coming soon
## Usage
Add this to your `Cargo.toml`:
```toml
[dependencies]
jit-assembler = "0.3"
```
### Basic Usage
```rust
use jit_assembler::riscv::{reg, csr, Riscv64InstructionBuilder};
use jit_assembler::riscv64_asm;
// Macro style (concise and assembly-like)
let instructions = riscv64_asm! {
csrrw(reg::RA, csr::MSTATUS, reg::SP); // CSR read-write using aliases
csrr(reg::T0, csr::MSTATUS); // CSR read (alias)
addi(reg::A0, reg::ZERO, 100); // Add immediate using aliases
add(reg::A1, reg::A0, reg::SP); // Register add with aliases
beq(reg::A0, reg::A1, 8); // Branch if equal
jal(reg::RA, 0x1000); // Jump and link
ret(); // Return (alias for jalr x0, x1, 0)
};
// Method chaining style (recommended for programmatic use)
let mut builder = Riscv64InstructionBuilder::new();
let instructions2 = builder
.csrrw(reg::RA, csr::MSTATUS, reg::SP) // CSR read-write using aliases
.addi(reg::A0, reg::ZERO, 100) // Add immediate with aliases
.add(reg::A1, reg::A0, reg::SP) // Register add using aliases
.beq(reg::A0, reg::A1, 8) // Branch if equal
.jal(reg::RA, 0x1000) // Jump and link
.ret() // Return instruction
.instructions();
// Traditional style
let mut builder3 = Riscv64InstructionBuilder::new();
builder3.csrrw(reg::RA, csr::MSTATUS, reg::SP);
builder3.addi(reg::A0, reg::ZERO, 100);
builder3.ret();
let instructions3 = builder3.instructions();
// Convert instructions to bytes easily
let bytes = instructions.to_bytes(); // All instructions as one byte vector
let size = instructions.total_size(); // Total size in bytes
let count = instructions.len(); // Number of instructions
// Iterate over instructions
for (i, instr) in instructions.iter().enumerate() {
println!("Instruction {}: {} -> {:?}", i, instr, instr.bytes());
}
// Or access by index
let first_instr = instructions[0];
```
### No-std Usage
For `no_std` environments, disable the default features:
```toml
[dependencies]
jit-assembler = { version = "0.3", default-features = false, features = ["riscv"] }
# Or for AArch64 only:
# jit-assembler = { version = "0.3", default-features = false, features = ["aarch64"] }
# Or for both architectures without std:
# jit-assembler = { version = "0.3", default-features = false, features = ["riscv", "aarch64"] }
```
## Architecture Support
### RISC-V
The RISC-V backend supports:
- **Base integer instruction set (RV64I)**:
- **Arithmetic**: `add`, `sub`, `addi`, `xor`, `or`, `and`, `slt`, `sltu`
- **Immediate arithmetic**: `andi`, `ori`, `xori`, `slti`, `sltiu`
- **Shifts**: `sll`, `srl`, `sra`, `slli`, `srli`, `srai`
- **Upper immediates**: `lui`, `auipc`
- **M extension (Integer Multiplication and Division)**:
- **Multiply**: `mul`, `mulh`, `mulhsu`, `mulhu`
- **Divide**: `div`, `divu`, `rem`, `remu`
- **Memory operations**:
- **Loads (signed)**: `ld`, `lw`, `lh`, `lb`
- **Loads (unsigned)**: `lbu`, `lhu`, `lwu`
- **Stores**: `sd`, `sw`, `sh`, `sb`
- **Control flow**: `jal`, `jalr`, `beq`, `bne`, `blt`, `bge`, `bltu`, `bgeu`
- **CSR instructions**: `csrrw`, `csrrs`, `csrrc`, `csrrwi`, `csrrsi`, `csrrci`
- **CSR pseudo-instructions**: `csrr` (read), `csrw` (write), `csrs` (set), `csrc` (clear), `csrwi`, `csrsi`, `csrci`
- **Privileged instructions**: `sret`, `mret`, `ecall`, `ebreak`, `wfi`
- **Pseudo-instructions**: `ret`, `li`
- **Register usage tracking**: Full tracking support for all instruction types (`register-tracking` feature)
### AArch64
The AArch64 backend supports:
- **Basic arithmetic operations**:
- **Register operations**: `add`, `sub`, `mul`, `udiv`, `sdiv`
- **Immediate operations**: `addi`, `subi`
- **Multiply-subtract operations**: `msub` (multiply-subtract for implementing remainder)
- **Logical operations**:
- **Register operations**: `and`, `or`, `xor` (EOR)
- **Move operations**: `mov`
- **Control flow**:
- **Return**: `ret`, `ret_reg`
- **Extended operations**:
- **Immediate moves**: `mov_imm` (for larger constants)
- **Shift operations**: `shl` (left shift using multiply)
- **Register conventions**: Following AAPCS64 (ARM ABI)
- **Argument/return registers**: X0-X7
- **Caller-saved temporaries**: X8-X18
- **Callee-saved registers**: X19-X28
- **Special registers**: X29 (FP), X30 (LR), X31 (SP/XZR)
- **Register usage tracking**: Full tracking support (`register-tracking` feature)
- **JIT compilation**: Direct function compilation and execution
### Future Architectures
Support for additional architectures is planned:
- x86-64: Intel/AMD 64-bit instruction set
## Examples
### JIT Compiler Integration
```rust
use jit_assembler::riscv::{reg, csr, Riscv64InstructionBuilder};
use jit_assembler::riscv64_asm;
// Simple function generator with macro
fn generate_add_function(a: i16, b: i16) -> Vec<u8> {
let instructions = riscv64_asm! {
addi(reg::A0, reg::ZERO, a); // Load first operand into a0
addi(reg::A1, reg::ZERO, b); // Load second operand into a1
add(reg::A0, reg::A0, reg::A1); // Add them, result in a0
ret(); // Return
};
// Convert to bytes for execution
instructions.to_bytes()
}
// Builder pattern for complex logic
fn generate_csr_routine() -> Vec<u8> {
let mut builder = Riscv64InstructionBuilder::new();
builder
.csrr(reg::T0, csr::MSTATUS) // Read current status into t0
.addi(reg::T1, reg::T0, 1) // Modify value in t1
.csrrw(reg::A0, csr::MSTATUS, reg::T1); // Write back, old value in a0
// Convert to executable code
builder.instructions().to_bytes()
}
```
### AArch64 Usage
```rust
use jit_assembler::aarch64::{reg, Aarch64InstructionBuilder};
use jit_assembler::common::InstructionBuilder;
use jit_assembler::aarch64_asm;
// Macro style (concise and assembly-like)
fn generate_aarch64_add_function_macro() -> Vec<u8> {
let instructions = aarch64_asm! {
add(reg::X0, reg::X0, reg::X1); // Add first two arguments (X0 + X1 -> X0)
ret(); // Return
};
instructions
}
// Builder pattern style
fn generate_aarch64_add_function() -> Vec<u8> {
let mut builder = Aarch64InstructionBuilder::new();
builder
.add(reg::X0, reg::X0, reg::X1) // Add first two arguments (X0 + X1 -> X0)
.ret(); // Return
builder.instructions().to_bytes()
}
// More complex AArch64 example with immediate values (macro style)
fn generate_aarch64_calculation_macro() -> Vec<u8> {
aarch64_asm! {
mov_imm(reg::X1, 42); // Load immediate 42 into X1
mul(reg::X0, reg::X0, reg::X1); // Multiply X0 by 42
addi(reg::X0, reg::X0, 100); // Add 100 to result
ret(); // Return
}
}
// More complex AArch64 example with immediate values (builder style)
fn generate_aarch64_calculation() -> Vec<u8> {
let mut builder = Aarch64InstructionBuilder::new();
builder
.mov_imm(reg::X1, 42) // Load immediate 42 into X1
.mul(reg::X0, reg::X0, reg::X1) // Multiply X0 by 42
.addi(reg::X0, reg::X0, 100) // Add 100 to result
.ret(); // Return
builder.instructions().to_bytes()
}
```
### JIT Execution (std-only)
Create and execute functions directly at runtime:
```rust
use jit_assembler::riscv64::{reg, Riscv64InstructionBuilder};
use jit_assembler::common::InstructionBuilder;
// Create a JIT function that adds two numbers
let add_func = unsafe {
Riscv64InstructionBuilder::new()
.add(reg::A0, reg::A0, reg::A1) // Add first two arguments
.ret() // Return result
.function::<fn(u64, u64) -> u64>()
}.expect("Failed to create JIT function");
// Call the JIT function naturally - just like a regular function!
let result = add_func.call(10, 20);
assert_eq!(result, 30);
// Create a function that returns a constant
let constant_func = unsafe {
Riscv64InstructionBuilder::new()
.addi(reg::A0, reg::ZERO, 42) // Load 42 into return register
.ret() // Return
.function::<fn() -> u64>()
}.expect("Failed to create JIT function");
let result = constant_func.call();
assert_eq!(result, 42);
```
**Note**: JIT execution requires the target architecture to match the host architecture. RISC-V code will only execute correctly on RISC-V systems.
**Features**:
- Type-safe function signatures
- Automatic memory management with `jit-allocator2`
- Natural function call syntax: `func.call()`, `func.call(arg)`, `func.call(arg1, arg2)`, etc. - just like regular functions!
- Cross-platform executable memory allocation
## Register Usage Tracking
The `register-tracking` feature enables comprehensive analysis of register usage patterns in your JIT-compiled code, helping with optimization and ABI compliance.
### Enable Register Tracking
Add the feature to your `Cargo.toml`:
```toml
[dependencies]
jit-assembler = { version = "0.3", features = ["register-tracking"] }
```
### Usage Example
```rust
use jit_assembler::riscv64::{reg, Riscv64InstructionBuilder};
use jit_assembler::common::InstructionBuilder;
let mut builder = Riscv64InstructionBuilder::new();
// Build a function that uses various registers
builder
.add(reg::T0, reg::T1, reg::T2) // T0 written, T1+T2 read
.addi(reg::T3, reg::SP, 16) // T3 written, SP read
.mul(reg::A0, reg::A1, reg::A2) // A0 written, A1+A2 read
.ld(reg::S0, reg::T0, 8) // S0 written, T0 read
.sd(reg::SP, reg::S1, -16); // SP+S1 read
// Analyze register usage
let usage = builder.register_usage();
println!("=== Register Usage Analysis ===");
println!("Total registers used: {}", usage.register_count());
println!("Written registers: {:?}", usage.written_registers());
println!("Read registers: {:?}", usage.read_registers());
// ABI compliance analysis
println!("Caller-saved (written): {:?}", usage.caller_saved_written());
println!("Callee-saved (written): {:?}", usage.callee_saved_written());
println!("Needs stack frame: {}", usage.needs_stack_frame());
// Detailed breakdown
let (caller, callee, special) = usage.count_by_abi_class();
println!("ABI breakdown - Caller: {}, Callee: {}, Special: {}",
caller, callee, special);
```
### Key Features
- **Separate tracking**: Distinguishes between written (def) and read (use) registers
- **ABI classification**: Automatically categorizes registers as caller-saved, callee-saved, or special-purpose
- **Stack frame analysis**: Determines if function prologue/epilogue is needed based on callee-saved register usage
- **Comprehensive coverage**: Tracks all RISC-V instruction types (R, I, S, B, U, J, CSR)
- **No-std compatible**: Uses `hashbrown` for no-std environments
### Register ABI Classification (RISC-V)
- **Caller-saved**: T0-T6, A0-A7, RA - Can be freely used without preservation
- **Callee-saved**: S0-S11, SP - Must be saved/restored if modified
- **Special**: X0 (zero), GP, TP - Require careful handling
This information is invaluable for:
- **Register allocation**: Choose optimal registers for variables
- **ABI compliance**: Ensure proper calling convention adherence
- **Performance optimization**: Minimize unnecessary register saves/restores
- **Code analysis**: Understand register pressure and usage patterns
## Merging Instruction Collections
When building complex functions, you may need to combine multiple instruction sequences in arbitrary order. For example, you might want to build the main function body first, then add prologue and epilogue code after determining which registers need to be saved.
### Basic Merging
```rust
use jit_assembler::riscv64::{reg, Riscv64InstructionBuilder};
use jit_assembler::common::InstructionBuilder;
// Build different parts separately
let mut prologue = Riscv64InstructionBuilder::new();
prologue
.addi(reg::SP, reg::SP, -16)
.sd(reg::SP, reg::RA, 8);
let mut main_code = Riscv64InstructionBuilder::new();
main_code
.add(reg::A0, reg::A1, reg::A2)
.mul(reg::A0, reg::A0, reg::A3);
let mut epilogue = Riscv64InstructionBuilder::new();
epilogue
.ld(reg::RA, reg::SP, 8)
.addi(reg::SP, reg::SP, 16)
.ret();
// Combine them using + operator: prologue + main + epilogue
let combined = prologue.instructions() + main_code.instructions() + epilogue.instructions();
// Or use method chaining
let mut combined = prologue.instructions();
combined += main_code.instructions();
combined += epilogue.instructions();
// Convert to executable code
let bytes = combined.to_bytes();
```
### Merging with Register Tracking
When the `register-tracking` feature is enabled, you can merge instruction collections while preserving register usage information:
```rust
use jit_assembler::riscv64::{reg, Riscv64InstructionBuilder};
use jit_assembler::common::{InstructionBuilder, InstructionCollectionWithUsage};
// Build code in separate builders
let mut prologue = Riscv64InstructionBuilder::new();
prologue.sd(reg::SP, reg::S0, -8); // Save S0
let mut main_code = Riscv64InstructionBuilder::new();
main_code.add(reg::S0, reg::A0, reg::A1); // Use S0
let mut epilogue = Riscv64InstructionBuilder::new();
epilogue.ld(reg::S0, reg::SP, -8); // Restore S0
// Create tracked collections
let prologue_tracked = InstructionCollectionWithUsage::new(
prologue.instructions(),
prologue.register_usage().clone()
);
let main_tracked = InstructionCollectionWithUsage::new(
main_code.instructions(),
main_code.register_usage().clone()
);
let epilogue_tracked = InstructionCollectionWithUsage::new(
epilogue.instructions(),
epilogue.register_usage().clone()
);
// Merge with register usage tracking
let combined = prologue_tracked + main_tracked + epilogue_tracked;
// Access both instructions and register usage
let instructions = combined.instructions();
let usage = combined.register_usage();
println!("Combined {} instructions", instructions.len());
println!("Register usage: {}", usage);
```
See `examples/instruction_collection_merge.rs` for a complete working example.
## Contributing
Contributions are welcome! Please feel free to submit a Pull Request.
## License
This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.