//! Closed binary image graphs, private-state placement, and symbol rewriting.

use std::mem;

use std::fmt;

use std::collections::dict;

use std::lang::alloc;

use std::lang::strings;

use std::lang::il;

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

/// Maximum bytes in an image's private state or one native data section.

export constant MAX_DATA: u32 = 4 * 1024 * 1024;

/// One reverse initializer edge used to promote address-bearing constants.

record Edge: Copy {

    /// Declaration containing the pointer initializer.

    owner: u32,

    /// Next reverse edge for the same target.

    next: ?u32,

}

/// Shared placement facts used by catalog generation and native emission.

export record Plan: Copy {

    /// Original data declaration indexes.

    dataIndex: dict::Dict,

    /// Original function declaration indexes.

    fnIndex: dict::Dict,

    /// Data declarations classified for the ordinary data layout routine.

    items: *mut [il::Data],

    /// Private byte offsets, indexed by the original data declaration.

    offsets: *mut [u32],

    /// Native layout order of the private declarations.

    order: *[data::DataSym],

    /// Number of private bytes.

    size: u32,

    /// Required private allocation alignment.

    alignment: u32,

    /// Entry function declaration index.

    entry: u32,

    /// The entry returns a status word.

    returnsStatus: bool,

}

/// Persistent names corresponding to the declaration indexes in Plan.

export record Names: Copy {

    /// Namespaced data declaration names.

    data: *[*[u8]],

    /// Namespaced function names and permitted native external names.

    functions: *[*[u8]],

}

/// Allocate through the binary codec's checked public arena interface.

export fn storage(arena: *mut alloc::Arena, size: u32, alignment: u32, count: u32) -> *mut [opaque] throws (il::binary::Error) {

    return try il::binary::storage(arena, size, alignment, count, 0);

}

/// Copy bytes whose lifetime would otherwise end with the decoded image.

export fn copy(arena: *mut alloc::Arena, bytes: *[u8]) -> *[u8] throws (il::binary::Error) {

    let result = try storage(arena, 1, 1, bytes.len) as *mut [u8];

    try! mem::copy(result, bytes);

    return result;

}

/// Intern persistent names without overflowing the caller's shared pool.

fn intern(pool: *mut strings::Pool, arena: *mut alloc::Arena, bytes: *[u8]) -> *[u8] throws (il::binary::Error) {

    if let existing = strings::find(pool, bytes) { return existing; }

    if pool.count >= pool.table.len / 2 { throw il::binary::error(0, "image string pool exhausted"); }

    return strings::intern(pool, try copy(arena, bytes));

}

/// Construct a trusted helper or namespace name with a decimal image ID.

export fn name(prefix: *[u8], id: u32, suffix: *[u8], pool: *mut strings::Pool, arena: *mut alloc::Arena) -> *[u8] throws (il::binary::Error) {

    let mut digits: [u8; 10] = undefined;

    let number = fmt::formatU32(id, &mut digits[..]);

    let size = prefix.len as u64 + number.len as u64 + suffix.len as u64;

    if size > 0x7FFFFFFF { throw il::binary::error(0, "image symbol is too long"); }

    let saved = alloc::save(arena);

    let bytes = try storage(arena, 1, 1, size as u32) as *mut [u8];

    let mut pos = try! mem::copy(bytes, prefix);

    set pos += try! mem::copy(&mut bytes[pos..], number);

    try! mem::copy(&mut bytes[pos..], suffix);

    if let existing = strings::find(pool, bytes) {

        alloc::restore(arena, saved);

        return existing;

    }

    if pool.count >= pool.table.len / 2 { throw il::binary::error(0, "image string pool exhausted"); }

    return strings::intern(pool, bytes);

}

/// Decode already established declaration references into an index.

export fn index(map: *dict::Dict, symbol: *[u8]) -> u32 throws (il::binary::Error) {

    let value = dict::get(map, symbol) else { throw il::binary::error(0, "unresolved image graph reference"); };

    return value as u32;

}

/// Width of one repetition, matching the shared data emitter.

export fn width(item: il::DataItem) -> u32 {

    match item {

        case il::DataItem::Val { typ, .. } => return il::typeSize(typ),

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

        case il::DataItem::Str(bytes) => return bytes.len,

        case il::DataItem::Undef => return 1,

    }

}

/// Classify transitive private pointers, then use the ordinary data layout order.

export fn plan(image: *il::binary::Image, arena: *mut alloc::Arena) -> Plan throws (il::binary::Error) {

    let program = &image.program;

    if program.data.len >= data::MAX_DATA_SYMS { throw il::binary::error(0, "too many image data symbols"); }

    if program.fns.len > 8192 { throw il::binary::error(0, "too many image functions"); }

    let mut dataIndex = try il::binary::dictionary(arena, program.data.len, 0);

    let mut fnIndex = try il::binary::dictionary(arena, program.fns.len, 0);

    let items = try storage(arena, @sizeOf(il::Data), @alignOf(il::Data), program.data.len) as *mut [il::Data];

    let offsets = try storage(arena, @sizeOf(u32), @alignOf(u32), items.len) as *mut [u32];

    let heads = try storage(arena, @sizeOf(?u32), @alignOf(?u32), items.len) as *mut [?u32];

    let queue = try storage(arena, @sizeOf(u32), @alignOf(u32), items.len) as *mut [u32];

    let mut queued: u32 = 0;

    let mut edgeCount: u64 = 0;

    for item, i in program.data {

        dict::insert(&mut dataIndex, item.name, i as i32);

        set items[i] = item;

        set offsets[i] = 0;

        set heads[i] = nil;

        if not item.readOnly { set queue[queued] = i; set queued += 1; }

        for value in item.values {

            if value.count == 0 { continue; }

            if let case il::DataItem::Sym(_) = value.item { set edgeCount += 1; }

        }

    }

    for function, i in program.fns {

        dict::insert(&mut fnIndex, function.name, i as i32);

        if function.isExtern and (function.name.len <= 11 or not mem::eq(&function.name[..11], "user::sys::")) {

            throw il::binary::error(0, "image extern must be a declared sys native function");

        }

    }

    if edgeCount > 0x7FFFFFFF { throw il::binary::error(0, "too many image initializer references"); }

    let edges = try storage(arena, @sizeOf(Edge), @alignOf(Edge), edgeCount as u32) as *mut [Edge];

    let mut edgeIndex: u32 = 0;

    for item, i in items {

        for value in item.values {

            if value.count == 0 { continue; }

            match value.item {

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

                    let target = try index(&dataIndex, symbol);

                    set edges[edgeIndex] = Edge { owner: i, next: heads[target] };

                    set heads[target] = edgeIndex;

                    set edgeIndex += 1;

                }

                case il::DataItem::Fn(symbol) => { let _ = try index(&fnIndex, symbol); }

                else => {},

            }

        }

    }

    let mut cursor: u32 = 0;

    while cursor < queued {

        let mut next = heads[queue[cursor]];

        while let e = next {

            let edge = edges[e];

            if items[edge.owner].readOnly {

                set items[edge.owner].readOnly = false;

                set queue[queued] = edge.owner;

                set queued += 1;

            }

            set next = edge.next;

        }

        set cursor += 1;

    }

    let mut alignment: u32 = 8;

    let mut expected: u64 = 0;

    for pass in 0..2 {

        for item in items {

            if item.readOnly or item.isZeroInit <> (pass == 1) { continue; }

            let a = item.alignment as u64;

            if a == 0 or a & (a - 1) <> 0 or a > MAX_DATA as u64 {

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

            }

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

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

            if expected > MAX_DATA as u64 { throw il::binary::error(0, "image private state exceeds capacity"); }

        }

    }

    let symbols = try storage(arena, @sizeOf(data::DataSym), @alignOf(data::DataSym), queued) as *mut [data::DataSym];

    let mut symbolCount: u32 = 0;

    let size = data::layoutSection(items, symbols, &mut symbolCount, 0, false);

    for symbol in symbols { set offsets[try index(&dataIndex, symbol.name)] = symbol.addr; }

    let entryName = image.entry else { throw il::binary::error(0, "image has no entry function"); };

    let entry = try index(&fnIndex, entryName);

    let function = program.fns[entry];

    if function.isExtern or function.params.len <> 1 or il::typeSize(function.params[0].type) <> 8 {

        throw il::binary::error(0, "image entry must receive one Env pointer");

    }

    let mut returnsStatus = false;

    let mut returnsVoid = false;

    for block in function.blocks {

        for instruction in block.instrs {

            if let case il::Instr::Ret { val } = instruction {

                if let value = val {

                    match value {

                        case il::Val::Undef => set returnsVoid = true,

                        else => set returnsStatus = true,

                    }

                } else { set returnsVoid = true; }

            }

        }

    }

    if returnsStatus and (returnsVoid or il::typeSize(function.returnType) <> 4) {

        throw il::binary::error(0, "image entry must return void or u32");

    }

    return Plan { dataIndex, fnIndex, items, offsets, order: symbols, size, alignment, entry, returnsStatus };

}

/// Allocate all names once; native externs alone retain their original spelling.

export fn names(program: *il::Program, id: u32, pool: *mut strings::Pool, arena: *mut alloc::Arena) -> Names throws (il::binary::Error) {

    let dataNames = try storage(arena, @sizeOf(*[u8]), @alignOf(*[u8]), program.data.len) as *mut [*[u8]];

    let fnNames = try storage(arena, @sizeOf(*[u8]), @alignOf(*[u8]), program.fns.len) as *mut [*[u8]];

    let prefix = try name("images::image_", id, "::", pool, arena);

    for item, i in program.data { set dataNames[i] = try qualified(prefix, item.name, pool, arena); }

    for function, i in program.fns {

        set fnNames[i] = try intern(pool, arena, function.name) if function.isExtern

            else try qualified(prefix, function.name, pool, arena);

    }

    return Names { data: dataNames, functions: fnNames };

}

/// Prefix an original whole-graph name, retaining its complete qualified path.

fn qualified(prefix: *[u8], suffix: *[u8], pool: *mut strings::Pool, arena: *mut alloc::Arena) -> *[u8] throws (il::binary::Error) {

    let size = prefix.len as u64 + suffix.len as u64;

    if size > 0x7FFFFFFF { throw il::binary::error(0, "image symbol is too long"); }

    let saved = alloc::save(arena);

    let bytes = try storage(arena, 1, 1, size as u32) as *mut [u8];

    let pos = try! mem::copy(bytes, prefix);

    try! mem::copy(&mut bytes[pos..], suffix);

    if let existing = strings::find(pool, bytes) { alloc::restore(arena, saved); return existing; }

    if pool.count >= pool.table.len / 2 { throw il::binary::error(0, "image string pool exhausted"); }

    return strings::intern(pool, bytes);

}

/// Rewrite a typed symbol reference, rejecting unresolved graph references.

fn val(value: il::Val, plan: *Plan, names: *Names) -> il::Val throws (il::binary::Error) {

    match value {

        case il::Val::DataSym(symbol) => return il::Val::DataSym(names.data[try index(&plan.dataIndex, symbol)]),

        case il::Val::FnAddr(symbol) => return il::Val::FnAddr(names.functions[try index(&plan.fnIndex, symbol)]),

        else => return value,

    }

}

/// Rewrite mutable block-edge arguments in their decoder-owned storage.

fn args(values: *mut [il::Val], plan: *Plan, names: *Names) throws (il::binary::Error) {

    for value, i in values { set values[i] = try val(value, plan, names); }

}

/// Copy immutable call arguments only when they contain a symbol to rename.

fn callArgs(values: *[il::Val], plan: *Plan, names: *Names, arena: *mut alloc::Arena) -> *[il::Val] throws (il::binary::Error) {

    for value in values {

        match value {

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

                let copied = try storage(arena, @sizeOf(il::Val), @alignOf(il::Val), values.len) as *mut [il::Val];

                for argument, i in values { set copied[i] = try val(argument, plan, names); }

                return copied;

            },

            else => {},

        }

    }

    return values;

}

/// Rewrite every symbolic operand without reconstructing any source-language types.

fn instruction(item: il::Instr, plan: *Plan, names: *Names, arena: *mut alloc::Arena) -> il::Instr throws (il::binary::Error) {

    match item {

        case il::Instr::Reserve { dst, size, alignment } => return il::Instr::Reserve { dst, size: try val(size, plan, names), alignment },

        case il::Instr::Store { typ, src, dst, offset } => return il::Instr::Store { typ, src: try val(src, plan, names), dst, offset },

        case il::Instr::Blit { dst, src, size, alignment } => return il::Instr::Blit { dst, src, size: try val(size, plan, names), alignment },

        case il::Instr::Copy { dst, val: value } => return il::Instr::Copy { dst, val: try val(value, plan, names) },

        case il::Instr::BinOp { op, typ, dst, a, b } => return il::Instr::BinOp { op, typ, dst, a: try val(a, plan, names), b: try val(b, plan, names) },

        case il::Instr::UnOp { op, typ, dst, a } => return il::Instr::UnOp { op, typ, dst, a: try val(a, plan, names) },

        case il::Instr::Zext { typ, dst, val: value } => return il::Instr::Zext { typ, dst, val: try val(value, plan, names) },

        case il::Instr::Sext { typ, dst, val: value } => return il::Instr::Sext { typ, dst, val: try val(value, plan, names) },

        case il::Instr::Call { retTy, dst, func, args: values } => {

            return il::Instr::Call { retTy, dst, func: try val(func, plan, names), args: try callArgs(values, plan, names, arena) };

        }

        case il::Instr::Ret { val: value } => {

            if let v = value { return il::Instr::Ret { val: try val(v, plan, names) }; }

            return item;

        }

        case il::Instr::Jmp { args: values, .. } => { try args(values, plan, names); return item; }

        case il::Instr::Br { op, typ, a, b, thenTarget, thenArgs, elseTarget, elseArgs } => {

            try args(thenArgs, plan, names); try args(elseArgs, plan, names);

            return il::Instr::Br { op, typ, a: try val(a, plan, names), b: try val(b, plan, names), thenTarget, thenArgs, elseTarget, elseArgs };

        }

        case il::Instr::Switch { val: value, defaultTarget, defaultArgs, cases } => {

            try args(defaultArgs, plan, names);

            for c in cases { try args(c.args, plan, names); }

            return il::Instr::Switch { val: try val(value, plan, names), defaultTarget, defaultArgs, cases };

        }

        case il::Instr::Ecall { dst, num, a0, a1, a2, a3 } => return il::Instr::Ecall {

            dst, num: try val(num, plan, names), a0: try val(a0, plan, names), a1: try val(a1, plan, names),

            a2: try val(a2, plan, names), a3: try val(a3, plan, names),

        },

        case il::Instr::Load { .. }, il::Instr::Sload { .. }, il::Instr::Unreachable,

             il::Instr::Ebreak, il::Instr::MemoryFence => return item,

    }

}

/// Rename a decoder-owned body in place; its storage survives until emission ends.

export fn function(function: *il::Fn, index: u32, plan: *Plan, names: *Names, arena: *mut alloc::Arena) -> il::Fn throws (il::binary::Error) {

    let mut result = *function;

    set result.name = names.functions[index];

    for block in result.blocks {

        for item, i in block.instrs { set block.instrs[i] = try instruction(item, plan, names, arena); }

    }

    return result;

}