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.1"

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.1", default-features = false, features = ["riscv"] }
# Or for AArch64 only:
# jit-assembler = { version = "0.1", default-features = false, features = ["aarch64"] }
# Or for both architectures without std:
# jit-assembler = { version = "0.1", 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.1", 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

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.