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:

[dependencies]
jit-assembler = "0.3"

Basic Usage

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:

[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

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

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:

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:

[dependencies]
jit-assembler = { version = "0.3", features = ["register-tracking"] }

Usage Example

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

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:

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 file for details.