//! Shared-backend image code, read-only initializers, and trusted linker glue.

use std::collections::dict;

use std::lang::alloc;

use std::lang::strings;

use std::lang::il;

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

use std::lang::gen::bitset;

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

use std::arch::rv64;

use std::arch::rv64::emit;

use std::arch::rv64::encode;

use super::graph;

/// One repeated private pointer initializer, expressed in state-relative bytes.

record Fixup: Copy {

    /// First pointer slot's private byte offset.

    offset: u32,

    /// Target declaration's private byte offset.

    target: u32,

    /// Number of consecutive eight-byte pointer slots.

    count: u32,

}

/// Persistent read-only definitions and scratch relocation runs.

export record Initializer: Copy {

    /// Shared initializer and immutable data definitions.

    items: *[il::Data],

    /// Private pointer runs applied to each new instance.

    fixups: *[Fixup],

}

/// Retain initializer payloads and keep ordinary RO/function relocations intact.

fn value(value: il::DataValue, plan: *graph::Plan, names: *graph::Names, arena: *mut alloc::Arena) -> il::DataValue throws (il::binary::Error) {

    let mut item = value.item;

    match item {

        case il::DataItem::Sym(symbol) => {

            let index = try graph::index(&plan.dataIndex, symbol);

            set item = il::DataItem::Sym(names.data[index]) if plan.items[index].readOnly

                else il::DataItem::Val { typ: il::Type::W64, val: plan.offsets[index] as i64 };

        }

        case il::DataItem::Fn(symbol) => set item = il::DataItem::Fn(names.functions[try graph::index(&plan.fnIndex, symbol)]),

        case il::DataItem::Str(bytes) => set item = il::DataItem::Str(try graph::copy(arena, bytes)),

        else => {},

    }

    return il::DataValue { item, count: 0 if graph::width(item) == 0 else value.count };

}

/// Flatten private declarations into one RO initializer using their native layout.

export fn initializer(plan: *graph::Plan, names: *graph::Names, stateName: *[u8], arena: *mut alloc::Arena, scratch: *mut alloc::Arena) -> Initializer throws (il::binary::Error) {

    let mut shared: u32 = 0;

    let mut count: u64 = plan.order.len as u64 * 2;

    let mut fixupCount: u64 = 0;

    for item in plan.items {

        if item.readOnly { set shared += 1; continue; }

        if item.isZeroInit { continue; }

        set count += item.values.len as u64;

        for v in item.values {

            if v.count == 0 { continue; }

            if let case il::DataItem::Sym(symbol) = v.item {

                if not plan.items[try graph::index(&plan.dataIndex, symbol)].readOnly { set fixupCount += 1; }

            }

        }

    }

    if count > 0x7FFFFFFF or fixupCount > 0x7FFFFFFF { throw il::binary::error(0, "too many image initializer values"); }

    let values = try graph::storage(arena, @sizeOf(il::DataValue), @alignOf(il::DataValue), count as u32) as *mut [il::DataValue];

    let fixups = try graph::storage(scratch, @sizeOf(Fixup), @alignOf(Fixup), fixupCount as u32) as *mut [Fixup];

    let items = try graph::storage(arena, @sizeOf(il::Data), @alignOf(il::Data), shared + 1) as *mut [il::Data];

    let mut valueCount: u32 = 0;

    let mut fixupIndex: u32 = 0;

    let mut offset: u32 = 0;

    for symbol in plan.order {

        let item = plan.items[try graph::index(&plan.dataIndex, symbol.name)];

        if symbol.addr > offset {

            set values[valueCount] = il::DataValue { item: il::DataItem::Undef, count: symbol.addr - offset };

            set valueCount += 1;

            set offset = symbol.addr;

        }

        if item.isZeroInit {

            set values[valueCount] = il::DataValue { item: il::DataItem::Undef, count: item.size };

            set valueCount += 1;

            set offset += item.size;

            continue;

        }

        for v in item.values {

            if v.count > 0 {

                if let case il::DataItem::Sym(target) = v.item {

                    let targetIndex = try graph::index(&plan.dataIndex, target);

                    if not plan.items[targetIndex].readOnly {

                        set fixups[fixupIndex] = Fixup { offset, target: plan.offsets[targetIndex], count: v.count };

                        set fixupIndex += 1;

                    }

                }

            }

            set values[valueCount] = try value(v, plan, names, arena);

            set valueCount += 1;

            set offset += graph::width(v.item) * v.count;

        }

    }

    set items[0] = il::Data { name: stateName, size: plan.size, alignment: plan.alignment,

        readOnly: true, isZeroInit: false, values: &values[..valueCount] };

    let mut itemIndex: u32 = 1;

    for item, i in plan.items {

        if not item.readOnly { continue; }

        let count = 1 if item.isZeroInit else item.values.len;

        let copied = try graph::storage(arena, @sizeOf(il::DataValue), @alignOf(il::DataValue), count) as *mut [il::DataValue];

        if item.isZeroInit { set copied[0] = il::DataValue { item: il::DataItem::Undef, count: item.size }; }

        else { for v, j in item.values { set copied[j] = try value(v, plan, names, arena); } }

        set items[itemIndex] = il::Data { name: names.data[i], size: item.size, alignment: item.alignment,

            readOnly: true, isZeroInit: false, values: copied };

        set itemIndex += 1;

    }

    return Initializer { items, fixups };

}

/// Reserve the caller's growable list using checked allocation before any append.

export fn appendData(target: *mut *mut [il::Data], items: *[il::Data], arena: *mut alloc::Arena) throws (il::binary::Error) {

    let total = (*target).len as u64 + items.len as u64;

    if total > data::MAX_DATA_SYMS as u64 { throw il::binary::error(0, "native data symbol capacity exceeded"); }

    let mut slice = *target;

    if total > slice.cap as u64 {

        let mut capacity: u32 = 16;

        while capacity < total as u32 { set capacity *= 2; }

        let grown = try graph::storage(arena, @sizeOf(il::Data), @alignOf(il::Data), capacity) as *mut [il::Data];

        for item, i in slice { set grown[i] = item; }

        set slice = @sliceOf(grown.ptr, slice.len, capacity);

    }

    let allocator = alloc::arenaAllocator(arena);

    for item in items { slice.append(item, allocator); }

    set *target = slice;

}

/// Check hard backend limits and a conservative expansion bound for one function.

fn bounds(e: *emit::Emitter, function: *il::Fn) -> u32 throws (il::binary::Error) {

    if function.params.len > rv64::ARG_REGS.len { throw il::binary::error(0, "image function has too many arguments"); }

    let mut blocks = function.blocks.len as u64 + 1;

    let mut words: u64 = 256;

    let mut reserve: u64 = 0;

    let mut operands: u64 = 0;

    let mut calls: u64 = 0;

    let mut branches: u64 = 0;

    for block in function.blocks {

        if block.params.len > 16 { throw il::binary::error(0, "image block has too many parameters"); }

        for item in block.instrs {

            set words += 64;

            set operands += 5;

            set branches += 2;

            match item {

                case il::Instr::Call { func, args, .. } => {

                    match func {

                        case il::Val::FnAddr(_), il::Val::Reg(_) => {},

                        else => throw il::binary::error(0, "invalid image native call target"),

                    }

                    if args.len > rv64::ARG_REGS.len { throw il::binary::error(0, "image call has too many arguments"); }

                    set words += args.len as u64 * 32;

                    set operands += args.len as u64;

                    set calls += 1;

                }

                case il::Instr::Jmp { args, .. } => {

                    set words += args.len as u64 * 32;

                    set operands += args.len as u64;

                }

                case il::Instr::Br { thenArgs, elseArgs, .. } => {

                    if thenArgs.len > 0 and elseArgs.len > 0 { throw il::binary::error(0, "image branch requires split argument edges"); }

                    set words += (thenArgs.len as u64 + elseArgs.len as u64) * 32;

                    set operands += thenArgs.len as u64 + elseArgs.len as u64;

                }

                case il::Instr::Switch { defaultArgs, cases, .. } => {

                    set words += defaultArgs.len as u64 * 32;

                    set operands += defaultArgs.len as u64;

                    for c in cases {

                        set words += 32 + c.args.len as u64 * 32;

                        set operands += c.args.len as u64;

                        set branches += 2;

                        if c.args.len > 0 { set blocks += 1; }

                    }

                }

                case il::Instr::Reserve { size, alignment, .. } => {

                    if alignment == 0 or alignment & (alignment - 1) <> 0 or alignment > graph::MAX_DATA {

                        throw il::binary::error(0, "image stack alignment exceeds capacity");

                    }

                    match size {

                        case il::Val::Imm(bytes) => {

                            if bytes < 0 or bytes > graph::MAX_DATA as i64 { throw il::binary::error(0, "image stack reservation exceeds capacity"); }

                            let a = alignment as u64;

                            set reserve = ((reserve + a - 1) & ~(a - 1)) + bytes as u64;

                            if reserve > graph::MAX_DATA as u64 { throw il::binary::error(0, "image stack frame exceeds capacity"); }

                        }

                        case il::Val::Reg(_) => {

                            if alignment > 2048 { throw il::binary::error(0, "dynamic image stack alignment exceeds backend capacity"); }

                        }

                        else => throw il::binary::error(0, "invalid image stack reservation"),

                    }

                }

                case il::Instr::Blit { size, alignment, .. } => {

                    let case il::Val::Imm(bytes) = size else { throw il::binary::error(0, "image blit requires an immediate size"); };

                    if bytes < 0 or bytes > graph::MAX_DATA as i64 or alignment == 0 or alignment & (alignment - 1) <> 0 {

                        throw il::binary::error(0, "image blit exceeds backend capacity");

                    }

                    // A spilled base can force an unrolled copy with adjusted addresses.

                    let width = il::typeSize(il::copyType(bytes as u32, alignment)) if bytes > 0 else 1;

                    set words += (bytes as u64 / width as u64 + 8) * 12;

                }

                else => {},

            }

        }

    }

    if blocks > e.labels.blockOffsets.len as u64 { throw il::binary::error(0, "image function exceeds block capacity"); }

    // Local jumps use one JAL slot and must stay strictly inside its positive range.

    if words >= 0x100000 / rv64::INSTR_SIZE as u64 { throw il::binary::error(0, "image function exceeds local branch reach"); }

    if words > (e.code.len - e.codeLen) as u64 { throw il::binary::error(0, "native code capacity exceeded"); }

    if e.labels.funcs.count >= e.labels.funcs.entries.len / 2 { throw il::binary::error(0, "native function symbol capacity exceeded"); }

    if dict::get(&e.labels.funcs, function.name) <> nil { throw il::binary::error(0, "duplicate native image function"); }

    if e.funcs.len >= e.funcs.cap or operands + e.pendingAddrLoads.len as u64 > e.pendingAddrLoads.cap as u64

        or calls + e.pendingCalls.len as u64 > e.pendingCalls.cap as u64

        or branches + e.pendingBranches.len as u64 > e.pendingBranches.cap as u64 {

        throw il::binary::error(0, "native image relocation capacity exceeded");

    }

    return reserve as u32;

}

/// Shared liveness state used to bound the spill-candidate buffer.

record Pressure: Copy {

    /// Registers live at the inspected instruction.

    live: bitset::Bitset,

    /// Number of live registers.

    count: u32,

}

/// Add one existing IL operand to the current live register set.

fn liveUse(reg: il::Reg, context: *mut opaque) {

    let pressure = context as *mut Pressure;

    if not bitset::contains(&pressure.live, reg.n) {

        bitset::put(&mut pressure.live, reg.n);

        set pressure.count += 1;

    }

}

/// Run the existing allocation stages, checking their fixed pressure capacity.

fn allocate(function: *il::Fn, arena: *mut alloc::Arena) -> regalloc::AllocResult throws (il::binary::Error) {

    let config = rv64::targetConfig();

    let live = try regalloc::liveness::analyze(function, arena)

        catch { throw il::binary::error(0, "image liveness arena exhausted"); };

    let saved = alloc::save(arena);

    let words = try graph::storage(arena, @sizeOf(u32), @alignOf(u32), bitset::wordsFor(live.maxReg)) as *mut [u32];

    let mut pressure = Pressure { live: bitset::init(words), count: 0 };

    for block, b in function.blocks {

        bitset::copy(&mut pressure.live, &live.liveOut[b]);

        set pressure.count = bitset::count(&pressure.live);

        if pressure.count > 256 { throw il::binary::error(0, "image exceeds live register capacity"); }

        let mut index = block.instrs.len;

        while index > 0 {

            set index -= 1;

            let item = block.instrs[index];

            if let dst = il::instrDst(item) {

                if bitset::contains(&pressure.live, dst.n) {

                    bitset::clear(&mut pressure.live, dst.n);

                    set pressure.count -= 1;

                }

            }

            il::forEachReg(item, liveUse, &mut pressure as *mut opaque);

            if pressure.count > 256 { throw il::binary::error(0, "image exceeds live register capacity"); }

        }

    }

    alloc::restore(arena, saved);

    let spill = try regalloc::spill::analyze(function, &live, config.allocatable.len, config.calleeSaved.len, config.slotSize, arena)

        catch { throw il::binary::error(0, "image spill arena exhausted"); };

    let assignment = try regalloc::assign::assign(function, &live, &spill, &config, arena)

        catch { throw il::binary::error(0, "image register assignment arena exhausted"); };

    return regalloc::AllocResult { assignments: assignment.assignments, spill, usedCalleeSaved: assignment.usedCalleeSaved };

}

/// Emit through the existing allocator and selector, propagating scratch exhaustion.

export fn function(generator: *mut rv64::Generator, function: *il::Fn, scratch: *mut alloc::Arena) throws (il::binary::Error) {

    let reserve = try bounds(&generator.e, function);

    let saved = alloc::save(scratch);

    let allocation = try allocate(function, scratch) catch error {

        alloc::restore(scratch, saved);

        throw error;

    };

    let frame = emit::computeFrame(allocation.spill.frameSize + reserve as i32, allocation.usedCalleeSaved,

        function.blocks.len, function.isLeaf, false);

    if frame.totalSize > 64 * 1024 {

        alloc::restore(scratch, saved);

        throw il::binary::error(0, "image function frame exceeds native stack");

    }

    rv64::isel::selectFn(&mut generator.e, &allocation, function);

    alloc::restore(scratch, saved);

}

/// Reserve a trusted leaf helper in the same global symbol and relocation tables.

fn helper(e: *mut emit::Emitter, name: *[u8], words: u32) throws (il::binary::Error) {

    if dict::get(&e.labels.funcs, name) <> nil { throw il::binary::error(0, "duplicate image helper symbol"); }

    if e.labels.funcs.count >= e.labels.funcs.entries.len / 2 { throw il::binary::error(0, "native function symbol capacity exceeded"); }

    if words > e.code.len - e.codeLen { throw il::binary::error(0, "native code capacity exceeded"); }

    if e.funcs.len >= e.funcs.cap or e.pendingCalls.len == e.pendingCalls.cap

        or e.pendingAddrLoads.len == e.pendingAddrLoads.cap {

        throw il::binary::error(0, "native image helper capacity exceeded");

    }

    emit::recordFunc(e, name);

    emit::recordFuncOffset(e, name);

}

/// Tail-call one explicitly bound primitive without limiting its physical address.

export fn trampoline(generator: *mut rv64::Generator, name: *[u8], address: u64) throws (il::binary::Error) {

    let e = &mut generator.e;

    try helper(e, name, 9);

    emit::loadImm(e, rv64::SCRATCH1, address as i64);

    emit::emit(e, encode::jalr(rv64::ZERO, rv64::SCRATCH1, 0));

}

/// Build trusted startup, address getters, and private-pointer relocation routines.

export fn helpers(generator: *mut rv64::Generator, id: u32, plan: *graph::Plan, names: *graph::Names,

    initial: *Initializer, pool: *mut strings::Pool, arena: *mut alloc::Arena) throws (il::binary::Error) {

    let stateName = initial.items[0].name;

    let entry = try graph::name("images::native_", id, "", pool, arena);

    let entryGetter = try graph::name("images::entry_", id, "", pool, arena);

    let initialGetter = try graph::name("images::initial_", id, "", pool, arena);

    let relocate = try graph::name("images::relocate_", id, "", pool, arena);

    let e = &mut generator.e;

    try helper(e, entry, 16);

    emit::emit(e, encode::mv(rv64::TP, rv64::A0));

    emit::emitLd(e, rv64::SP, rv64::A0, 72);

    emit::recordCall(e, names.functions[plan.entry]);

    if not plan.returnsStatus { emit::emit(e, encode::mv(rv64::A0, rv64::ZERO)); }

    else {

        emit::emit(e, encode::slli(rv64::A0, rv64::A0, 32));

        emit::emit(e, encode::srli(rv64::A0, rv64::A0, 32));

    }

    emit::loadImm(e, rv64::A7, 49);

    emit::emit(e, encode::ecall());

    emit::emit(e, encode::ebreak());

    try helper(e, entryGetter, 3);

    emit::recordAddrLoad(e, entry, rv64::A0);

    emit::emit(e, encode::ret());

    try helper(e, initialGetter, 3);

    emit::recordDataAddrLoad(e, stateName, rv64::A0);

    emit::emit(e, encode::ret());

    let words = initial.fixups.len as u64 * 32 + 1;

    if words > 0x7FFFFFFF { throw il::binary::error(0, "image relocation helper exceeds capacity"); }

    try helper(e, relocate, words as u32);

    for fixup in initial.fixups {

        emit::emitAddImm(e, rv64::T0, rv64::A0, fixup.offset as i32);

        emit::emitAddImm(e, rv64::T1, rv64::A0, fixup.target as i32);

        emit::loadImm(e, rv64::T2, fixup.count as i64);

        let start = e.codeLen;

        if fixup.offset % 8 == 0 { emit::emitSd(e, rv64::T1, rv64::T0, 0); }

        else {

            emit::emit(e, encode::mv(rv64::T3, rv64::T1));

            for byte in 0..8 {

                emit::emitSb(e, rv64::T3, rv64::T0, byte as i32);

                if byte < 7 { emit::emit(e, encode::srli(rv64::T3, rv64::T3, 8)); }

            }

        }

        if fixup.count > 1 {

            emit::emit(e, encode::addi(rv64::T0, rv64::T0, 8));

            emit::emit(e, encode::addi(rv64::T2, rv64::T2, -1));

            emit::emit(e, encode::bne(rv64::T2, rv64::ZERO, (start as i32 - e.codeLen as i32) * rv64::INSTR_SIZE));

        }

    }

    emit::emit(e, encode::ret());

}