//! Position-independent native layout and persistent pointer initialization.

use std::mem;

use std::collections::dict;

use std::lang::alloc;

use std::lang::il;

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

use std::arch::rv64;

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

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

use super::graph;

/// Native bytes and offsets, with data addresses relative to the allocation.

export record Layout: Copy {

    /// Fully initialized read-only bytes with allocation-relative pointers.

    bytes: *[u8],

    /// Allocation-relative immutable symbol addresses.

    symbols: data::DataSymMap,

    /// Byte offset of the first native instruction.

    codeOffset: u32,

    /// Complete allocation size including native instructions.

    size: u32,

    /// Required power-of-two allocation alignment.

    alignment: u32,

}

/// Resolve a function before invoking backend relocation routines.

export fn functionOffset(e: *emit::Emitter, name: *[u8]) -> u32 throws (il::binary::Error) {

    let offset = dict::get(&e.labels.funcs, name) else {

        throw il::binary::error(0, "unresolved native image function");

    };

    if offset < 0 or offset as u32 % 4 <> 0 or offset as u32 / 4 >= e.codeLen {

        throw il::binary::error(0, "native image function offset exceeds code");

    }

    return offset as u32;

}

/// Require every reserved relocation instruction to belong to emitted code.

fn slots(e: *emit::Emitter, index: u32, count: u32) throws (il::binary::Error) {

    if index > e.codeLen or count > e.codeLen - index {

        throw il::binary::error(0, "native image relocation exceeds code");

    }

}

/// Check all references before the shared emitter's infallible final patching.

fn references(e: *emit::Emitter, items: *[il::Data], symbols: *data::DataSymMap) throws (il::binary::Error) {

    if e.pendingBranches.len <> 0 { throw il::binary::error(0, "unresolved native image block branch"); }

    for pending in e.pendingCalls {

        try slots(e, pending.index, 2);

        let _ = try functionOffset(e, pending.target);

    }

    for pending in e.pendingJumps {

        try slots(e, pending.index, 1);

        let offset = try functionOffset(e, pending.target);

        if not encode::isJumpImm(offset as i32 - pending.index as i32 * rv64::INSTR_SIZE) {

            throw il::binary::error(0, "native image jump exceeds reach");

        }

    }

    for pending in e.pendingAddrLoads {

        try slots(e, pending.index, 2);

        if pending.isData {

            if data::lookupAddr(symbols, pending.target) == nil {

                throw il::binary::error(0, "unresolved native image data address");

            }

        } else { let _ = try functionOffset(e, pending.target); }

    }

    for item in items {

        for value in item.values {

            if value.count == 0 { continue; }

            match value.item {

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

                    if data::lookupAddr(symbols, name) == nil {

                        throw il::binary::error(0, "unresolved native image data initializer");

                    }

                }

                case il::DataItem::Fn(name) => { let _ = try functionOffset(e, name); }

                else => {},

            }

        }

    }

}

/// Lay out flattened immutable data using the shared data placement and emitter.

/// Code and data address instructions use allocation-local PC-relative offsets.

export fn prepare(generator: *mut rv64::Generator, items: *[il::Data], arena: *mut alloc::Arena,

    scratch: *mut alloc::Arena) -> Layout throws (il::binary::Error) {

    let mut alignment: u32 = 8;

    let mut extent: u64 = 0;

    for item in items {

        let a = item.alignment as u64;

        if not item.readOnly or item.isZeroInit or a == 0 or a & (a - 1) <> 0 or a > graph::MAX_DATA as u64 {

            throw il::binary::error(0, "invalid flattened native image data");

        }

        if item.alignment > alignment { set alignment = item.alignment; }

        set extent = ((extent + a - 1) & ~(a - 1)) + item.size as u64;

        if extent > graph::MAX_DATA as u64 { throw il::binary::error(0, "native image data exceeds capacity"); }

        let mut size: u64 = 0;

        for value in item.values {

            set size += graph::width(value.item) as u64 * value.count as u64;

            if size > item.size as u64 { throw il::binary::error(0, "native image initializer exceeds data"); }

        }

        if size <> item.size as u64 { throw il::binary::error(0, "native image initializer size mismatch"); }

    }

    let syms = try graph::storage(scratch, @sizeOf(data::DataSym), @alignOf(data::DataSym), items.len) as *mut [data::DataSym];

    let mut count: u32 = 0;

    let dataSize = data::layoutSection(items, syms, &mut count, 0, true);

    let mut map = try il::binary::dictionary(scratch, count, 0);

    for symbol in syms {

        if dict::get(&map, symbol.name) <> nil { throw il::binary::error(0, "duplicate native image data symbol"); }

        dict::insert(&mut map, symbol.name, symbol.addr as i32);

    }

    let symbols = data::DataSymMap { dict: map, syms };

    let e = &mut generator.e;

    let codeOffset = mem::alignUp(dataSize, 8);

    let total = codeOffset as u64 + e.codeLen as u64 * rv64::INSTR_SIZE as u64;

    // Keep every allocation-local AUIPC displacement strictly inside signed reach.

    if total > 0x7FFFF000 { throw il::binary::error(0, "native image exceeds relative address reach"); }

    try references(e, items, &symbols);

    let bytes = try graph::storage(arena, 1, 1, dataSize) as *mut [u8];

    for i in 0..bytes.len { set bytes[i] = 0; }

    emit::patchCalls(e);

    emit::patchJumps(e);

    for pending in e.pendingAddrLoads {

        let mut target: u32 = 0;

        if pending.isData {

            let address = data::lookupAddr(&symbols, pending.target) else {

                throw il::binary::error(0, "unresolved native image data address");

            };

            set target = address;

        } else { set target = codeOffset + (try functionOffset(e, pending.target)); }

        let pc = codeOffset + pending.index * rv64::INSTR_SIZE as u32;

        let split = emit::splitImm(target as i32 - pc as i32);

        emit::patch(e, pending.index, encode::auipc(pending.rd, split.hi));

        emit::patch(e, pending.index + 1, encode::addi(pending.rd, pending.rd, split.lo));

    }

    data::emitSection(items, &symbols, &e.labels, codeOffset, bytes, true);

    return Layout { bytes, symbols, codeOffset, size: total as u32, alignment };

}

/// Convert allocation-relative initializer pointers to actual native addresses.

/// Scalar private offsets remain unchanged for the per-instance relocation helper.

export fn rebase(items: *[il::Data], destination: *mut [u8], base: u64) {

    let mut offset: u32 = 0;

    for item in items {

        set offset = mem::alignUp(offset, item.alignment);

        for value in item.values {

            match value.item {

                case il::DataItem::Sym(_), il::DataItem::Fn(_) => {

                    for _ in 0..value.count {

                        let mut relative: u64 = 0;

                        for byte in 0..8 { set relative |= (destination[offset + byte] as u64) << (byte as u64 * 8); }

                        let address = base + relative;

                        for byte in 0..8 { set destination[offset + byte] = (address >> (byte as u64 * 8)) as u8; }

                        set offset += 8;

                    }

                }

                else => set offset += graph::width(value.item) * value.count,

            }

        }

    }

}