//! RV64 code generation backend.

//!

//! Generates RISC-V 64-bit machine code from IL (intermediate language).

//!

//! # Submodules

//!

//! * encode: Instruction encoding functions

//! * decode: Instruction decoding (for disassembly/printing)

//! * emit: Binary emission context and branch patching

//! * isel: Instruction selection (IL to RV64 instructions)

//! * printer: Assembly text output

export mod encode;

export mod decode;

export mod emit;

export mod isel;

export mod printer;

export mod asm;

@test mod tests;

use std::mem;

use std::collections::dict;

use std::lang::il;

use std::lang::alloc;

use std::lang::gen;

use std::lang::gen::labels;

use std::lang::gen::regalloc;

use std::lang::gen::data;

use std::lang::gen::types;

/// Recoverable code-generation failure.

export union Error: Copy {

    /// The code-generation or function arena is full.

    Allocation,

    /// A fixed output or metadata table is full.

    Capacity,

    /// A required symbol is missing or invalid.

    Symbol,

    /// A branch or address exceeds its instruction range.

    Relocation,

    /// Data section layout or emission failed.

    Data(data::Error),

}

////////////////

// Registers  //

////////////////

export constant ZERO: gen::Reg = gen::Reg(0);   /// Hard-wired zero.

export constant RA:   gen::Reg = gen::Reg(1);   /// Return address.

export constant SP:   gen::Reg = gen::Reg(2);   /// Stack pointer.

export constant GP:   gen::Reg = gen::Reg(3);   /// Global pointer.

export constant TP:   gen::Reg = gen::Reg(4);   /// Thread pointer.

export constant T0:   gen::Reg = gen::Reg(5);   /// Temporary/alternate link register.

export constant T1:   gen::Reg = gen::Reg(6);   /// Temporary.

export constant T2:   gen::Reg = gen::Reg(7);   /// Temporary.

export constant S0:   gen::Reg = gen::Reg(8);   /// Saved register/frame pointer.

export constant FP:   gen::Reg = gen::Reg(8);   /// Frame pointer (alias for S0).

export constant S1:   gen::Reg = gen::Reg(9);   /// Saved register.

export constant A0:   gen::Reg = gen::Reg(10);  /// Function argument/return.

export constant A1:   gen::Reg = gen::Reg(11);  /// Function argument/return.

export constant A2:   gen::Reg = gen::Reg(12);  /// Function argument.

export constant A3:   gen::Reg = gen::Reg(13);  /// Function argument.

export constant A4:   gen::Reg = gen::Reg(14);  /// Function argument.

export constant A5:   gen::Reg = gen::Reg(15);  /// Function argument.

export constant A6:   gen::Reg = gen::Reg(16);  /// Function argument.

export constant A7:   gen::Reg = gen::Reg(17);  /// Function argument.

export constant S2:   gen::Reg = gen::Reg(18);  /// Saved register.

export constant S3:   gen::Reg = gen::Reg(19);  /// Saved register.

export constant S4:   gen::Reg = gen::Reg(20);  /// Saved register.

export constant S5:   gen::Reg = gen::Reg(21);  /// Saved register.

export constant S6:   gen::Reg = gen::Reg(22);  /// Saved register.

export constant S7:   gen::Reg = gen::Reg(23);  /// Saved register.

export constant S8:   gen::Reg = gen::Reg(24);  /// Saved register.

export constant S9:   gen::Reg = gen::Reg(25);  /// Saved register.

export constant S10:  gen::Reg = gen::Reg(26);  /// Saved register.

export constant S11:  gen::Reg = gen::Reg(27);  /// Saved register.

export constant T3:   gen::Reg = gen::Reg(28);  /// Temporary.

export constant T4:   gen::Reg = gen::Reg(29);  /// Temporary.

export constant T5:   gen::Reg = gen::Reg(30);  /// Temporary.

export constant T6:   gen::Reg = gen::Reg(31);  /// Temporary.

/// Create a register from a number. Panics if `n > 31`.

export fn reg(n: u8) -> gen::Reg {

    assert n < 32;

    return gen::Reg(n);

}

////////////////////////////

// Architecture constants //

////////////////////////////

/// Total number of general-purpose registers.

export constant NUM_REGISTERS: u8 = 32;

/// Number of saved registers.

export constant NUM_SAVED_REGISTERS: u8 = 11;

/// Word size in bytes (32-bit).

export constant WORD_SIZE: i32 = 4;

/// Doubleword size in bytes (64-bit).

export constant DWORD_SIZE: i32 = 8;

/// Instruction size in bytes.

export constant INSTR_SIZE: i32 = 4;

/// Stack alignment requirement in bytes.

export constant STACK_ALIGNMENT: i32 = 16;

/// Minimum blit size (in bytes) to use a loop instead of inline copy.

/// Blits below this threshold are fully unrolled as LD/SD pairs.

export constant BLIT_LOOP_THRESHOLD: i32 = 256;

/////////////////////////

// Codegen Allocation  //

/////////////////////////

/// Argument registers for function calls.

export constant ARG_REGS: [gen::Reg; 8] = [A0, A1, A2, A3, A4, A5, A6, A7];

/// Scratch register for code gen. Never allocated to user values.

export constant SCRATCH1: gen::Reg = T5;

/// Second scratch register for operations needing two temporaries.

export constant SCRATCH2: gen::Reg = T6;

/// Dedicated scratch for address offset adjustment. Never allocated to user

/// values and never used for operand materialization, so it can never

/// conflict with `rd`, `rs`, or `base` in load/store helpers.

export constant ADDR_SCRATCH: gen::Reg = T4;

/// Callee-saved registers that need save/restore if used.

export constant CALLEE_SAVED: [gen::Reg; NUM_SAVED_REGISTERS] = [S1, S2, S3, S4, S5, S6, S7, S8, S9, S10, S11];

/// Maximum 12-bit signed immediate value.

export constant MAX_IMM: i32 = 2047;

/// Minimum 12-bit signed immediate value.

export constant MIN_IMM: i32 = -2048;

/// Allocatable registers for register allocation.

constant ALLOCATABLE_REGS: [gen::Reg; 23] = [

    T0, T1, T2, T3,                             // Temporaries

    A0, A1, A2, A3, A4, A5, A6, A7,             // Arguments

    S1, S2, S3, S4, S5, S6, S7, S8, S9, S10, S11, // Saved

];

/// Get target configuration for register allocation.

// TODO: This should be a constant variable.

export fn targetConfig() -> regalloc::TargetConfig {

    return regalloc::TargetConfig {

        allocatable: &ALLOCATABLE_REGS[..],

        argRegs: &ARG_REGS[..],

        calleeSaved: &CALLEE_SAVED[..],

        slotSize: DWORD_SIZE,

    };

}

///////////////////////

// Codegen Constants //

///////////////////////

/// Base address where read-only data is loaded.

export constant RO_DATA_BASE: u32 = 0x10000;

/// Base address where read-write data is loaded.

export constant RW_DATA_BASE: u32 = 0xFFFFF0;

/// Single-file RV64 image magic, "RAD0" as a little-endian u32.

export constant IMAGE_MAGIC: u32 = 0x30444152;

/// Single-file RV64 image format version.

export constant IMAGE_VERSION: u32 = 1;

/// Build a single-file RV64 image header.

export fn imageHeader(codeBytes: u32, roDataBytes: u32, rwDataBytes: u32) -> [u32; 5] {

    return [IMAGE_MAGIC, IMAGE_VERSION, codeBytes, roDataBytes, rwDataBytes];

}

/// Storage buffers passed from driver for code generation.

export record Storage {

    /// Buffer for data symbols.

    dataSyms: *mut [data::DataSym],

    /// Hash table entries for data symbol lookup.

    dataSymEntries: *mut [dict::Entry],

}

/// Result of code generation.

export record Program: Copy {

    /// Slice of emitted code.

    code: *[u32],

    /// Slice of function addresses (name + start index).

    funcs: *[types::FuncAddr],

    /// Number of read-only data bytes emitted.

    roDataSize: u32,

    /// Number of read-write data bytes emitted.

    rwDataSize: u32,

    /// Debug entries mapping PCs to source locations. Empty when debug is off.

    debugEntries: *[types::DebugEntry],

}

/// Entry jump patching requested for the generated program.

export union EntryPatch: Copy {

    /// No entry jump is emitted.

    None,

    /// Reserve code slot zero for a jump to the default function.

    Reserved(?*[u8]),

}

/// Options controlling incremental RV64 program generation.

export record ProgramOptions: Copy {

    /// Entry jump patching mode.

    entryPatch: EntryPatch,

    /// Whether to emit debug source locations.

    debug: bool,

}

/// State for incremental RV64 program generation.

///

/// The generator owns global codegen state that must survive across function

/// emission. Function-local scratch stays outside this record so callers can

/// reclaim it after each function.

export record Generator {

    /// Binary emitter and relocation state.

    e: emit::Emitter,

    /// Entry jump patching state.

    entryPatch: EntryPatch,

}

/// Begin RV64 code generation for a program's global state.

export unsafe fn beginProgram(

    options: ProgramOptions,

    arena: &mut alloc::Arena

) -> Generator throws (Error) {

    let checkpoint = alloc::save(arena);

    let mut e = try emit::emitter(arena, options.debug) catch {

        alloc::restore(arena, checkpoint);

        throw Error::Allocation;

    };

    // Emit placeholder entry jump when requested.

    // We'll patch this at the end once we know where the function is.

    match options.entryPatch {

        case EntryPatch::Reserved(_) => {

            emit::emit(&mut e, encode::nop()); // Placeholder for two-instruction jump.

            emit::emit(&mut e, encode::nop()); //

        }

        else => {}

    }

    return Generator {

        e,

        entryPatch: options.entryPatch,

    };

}

/// Generate code for one IL function.

export unsafe fn generateFunction(

    generator: &mut Generator,

    func: *unsafe il::Fn,

    arena: &mut alloc::Arena

) {

    if generator.e.error <> nil or func.isExtern {

        return;

    }

    let checkpoint = alloc::save(arena);

    let config = targetConfig();

    let ralloc = try regalloc::allocate(func, &config, arena) catch {

        alloc::restore(arena, checkpoint);

        set generator.e.error = Error::Allocation;

        return;

    };

    isel::selectFn(&mut generator.e, &ralloc, func);

    // Reclaim unused memory after instruction selection.

    alloc::restore(arena, checkpoint);

}

/// Record an alternate name for the next function emitted.

export fn recordFunctionAlias(generator: &mut Generator, name: *[u8]) {

    let codeLen = generator.e.codeLen;

    emit::recordFuncOffsetAt(&mut generator.e, name, codeLen);

}

/// Add the text section of an assembled program to the generator.

///

/// This function snapshots the generator's current code length as the base

/// index, converts each text symbol's byte offset to an instruction index, adds

/// that base, and records the final address for printing. Only `.export` text

/// symbols are exported to the emitter's function-offset table for extern call

/// resolution. Local labels must not escape their assembly fragment because

/// separate assembly inputs may reuse the same local names.

///

/// Non-text symbols are ignored here because assembled data is not appended to

/// the generator's text stream. The driver merges assembled data into the RO data

/// prefix separately and passes that data to [`finishProgram`].

export fn addAssembly(generator: &mut Generator, program: asm::Program) {

    let baseIndex = generator.e.codeLen;

    for symbol in program.symbols {

        if symbol.section == asm::Section::Text {

            let index = baseIndex + ((symbol.offset as u32) / INSTR_SIZE as u32);

            emit::recordFuncAt(&mut generator.e, symbol.name, index);

            if symbol.isExported {

                emit::recordFuncOffsetAt(&mut generator.e, symbol.name, index);

            }

        }

    }

    for fixup in program.externalFixups {

        match fixup.info {

            case asm::FixupInfo::Jal { rd, index } => {

                emit::recordJumpAt(&mut generator.e, fixup.symbol, rd, baseIndex + index);

            }

            case asm::FixupInfo::Addr { rd, index } => {

                emit::recordAddrLoadAt(&mut generator.e, fixup.symbol, rd, baseIndex + index);

            }

            else => panic "addAssembly: invalid external fixup",

        }

    }

    for word in program.text {

        emit::emit(&mut generator.e, word);

    }

}

/// Finish RV64 code generation and return the emitted program.

export unsafe fn finishProgram(

    generator: &mut Generator,

    globalData: &[il::Data],

    storage: Storage,

    roDataPrefix: *[u8],

    roDataBuf: &mut [u8],

    rwDataBuf: &mut [u8]

) -> Program throws (Error) {

    try emit::check(&generator.e);

    // Build data map after function lowering. Function-local literals can add

    // global data while functions are lowered, so final layout belongs here.

    let case Storage { dataSyms: symbolBuf, dataSymEntries } = storage

        else panic "expected code generation storage";

    let mut dataSymCount: u32 = 0;

    let roLayoutSize = try data::layoutSectionAtOffset(

        globalData, symbolBuf, &mut dataSymCount, RO_DATA_BASE, roDataPrefix.len, true

    ) catch error {

        throw Error::Data(error);

    };

    try data::layoutSection(globalData, symbolBuf, &mut dataSymCount, RW_DATA_BASE, false) catch error {

        throw Error::Data(error);

    };

    let dataSyms = &symbolBuf[..dataSymCount];

    let dataSymMap = try data::buildMap(dataSyms, dataSymEntries) catch error {

        throw Error::Data(error);

    };

    let codeBase64 = (RO_DATA_BASE as u64 + roLayoutSize as u64 + DWORD_SIZE as u64 - 1)

        & ~(DWORD_SIZE as u64 - 1);

    if codeBase64 > 0xffffffff {

        throw Error::Data(data::Error::Overflow);

    }

    let codeBase = codeBase64 as u32;

    match generator.entryPatch {

        case EntryPatch::Reserved(targetName) => {

            let target = targetName else {

                throw Error::Symbol;

            };

            let offset = emit::branchOffsetToFunc(&mut generator.e, 0, target);

            let s = emit::splitImm(offset);

            emit::patch(&mut generator.e, 0, encode::auipc(SCRATCH1, s.hi));

            emit::patch(&mut generator.e, 1, encode::jalr(ZERO, SCRATCH1, s.lo));

        }

        else => {}

    }

    // Patch function calls and address loads now that all functions are emitted.

    emit::patchJumps(&mut generator.e);

    emit::patchCalls(&mut generator.e);

    emit::patchAddrLoads(&mut generator.e, &dataSymMap);

    try emit::check(&generator.e);

    // Emit data sections.

    if roDataPrefix.len > roDataBuf.len {

        throw Error::Capacity;

    }

    try! mem::copy(roDataBuf, roDataPrefix);

    let roDataSize = try data::emitSectionAtOffset(

        globalData, &dataSymMap, &generator.e.labels, codeBase, roDataBuf, true, roDataPrefix.len

    ) catch error {

        throw Error::Data(error);

    };

    let rwDataSize = try data::emitSection(

        globalData, &dataSymMap, &generator.e.labels, codeBase, rwDataBuf, false

    ) catch error {

        throw Error::Data(error);

    };

    return Program {

        code: emit::getCode(&generator.e),

        funcs: &generator.e.funcs[..],

        roDataSize,

        rwDataSize,

        debugEntries: emit::getDebugEntries(&generator.e),

    };

}