// Expect retcode = 0

// 目标布局（移植自 stage1/spl_layout.c）。输入是 IR 类型表 TypeTable；
// 为每个类型计算 size_bits / align_bits / 字段位偏移，供 ir2vm 使用。
// 默认 64 位平台：标量按 1/2/4/8 字节对齐，指针 = usize，聚合按最大字段对齐
// 且不允许重排（与 stage1 的默认策略一致）。

type .{ .Error } = @import("utils/error.spl");
type .{ .Vec } = @import("utils/vec.spl");
type TypeTable = @import("./type_table.spl");
type .{ .TypeNode, .Agg, .AggKind, .Member } = @import("./type_table.spl");
type .{ .node, .member_at, .make_int, .make_array, .make_agg, .push_member } = @import("./type_table.spl");

type LayoutInfo = struct {
    size_bits: usize,
    align_bits: usize,
    field_bit_offsets: Vec, // usize
    field_bit_sizes: Vec,   // usize
    resolved: bool,
}

type Layout = @self;

infos: Vec, // LayoutInfo
ptr_size: usize,
ptr_align: usize,

type LayoutResult = enum {
    Ok: Layout,
    Err: Error,
}

fn blank_usizes(): Vec {
    ret Vec.blank(@sizeof(usize), @alignof(usize));
}

fn new(): LayoutResult {
    ret .Ok(Layout {
        .infos = Vec.blank(@sizeof(LayoutInfo), @alignof(LayoutInfo)),
        .ptr_size = @sizeof(usize),
        .ptr_align = @alignof(usize),
    });
}

fn info_at(self: *Layout, id: usize): *LayoutInfo {
    if id == 0 || id >= self.infos.len { ret null; }
    ret (self.infos.data as *LayoutInfo) + id;
}

// 未解析或越界返回 null（对应 stage1 spl_layout_info_of）。
fn info_of(self: *Layout, id: usize): *LayoutInfo {
    var info = info_at(self, id);
    if info == null || !info.resolved { ret null; }
    ret info;
}

fn drop(self: *Layout) {
    var i: usize = 0;
    while i < self.infos.len {
        var info = info_at(self, i);
        if info != null {
            Vec.drop(&info.field_bit_offsets);
            Vec.drop(&info.field_bit_sizes);
        }
        i = i + 1;
    }
    Vec.drop(&self.infos);
}

fn ensure(self: *Layout, n: usize) {
    while self.infos.len < n {
        var info = LayoutInfo {
            .size_bits = 0,
            .align_bits = 0,
            .field_bit_offsets = blank_usizes(),
            .field_bit_sizes = blank_usizes(),
            .resolved = false,
        };
        match Vec.push(&self.infos, &info as *u8) {
            .Ok => {},
            .Err(error) => { ret; },
        }
    }
}

fn push_usize(v: *Vec, x: usize) {
    match Vec.push(v, &x as *u8) {
        .Ok => {},
        .Err(error) => {},
    }
}

fn usize_at(v: *Vec, i: usize): usize {
    match Vec.get(v, i) {
        .Ok(ptr) => { ret (ptr as *usize).*; },
        .Err(error) => { ret 0; },
    }
    ret 0;
}

// 标量对齐（字节）；0 表示没有匹配的标量（即未落地）。
fn scalar_align(size_bytes: usize): usize {
    if size_bytes == 1 { ret 1; }
    if size_bytes == 2 { ret 2; }
    if size_bytes == 4 { ret 4; }
    if size_bytes == 8 { ret 8; }
    ret 0;
}

fn round_up_bits(x: usize, align_bits: usize): usize {
    if align_bits <= 1 { ret x; }
    ret (x + align_bits - 1) / align_bits * align_bits;
}

fn layout_scalar(info: *LayoutInfo, size_bits: usize) {
    var size_bytes = (size_bits + 7) / 8;
    var a = scalar_align(size_bytes);
    if a == 0 { ret; }
    info.size_bits = size_bytes * 8;
    info.align_bits = a * 8;
}

fn layout_ptr(self: *Layout, info: *LayoutInfo) {
    if self.ptr_size == 0 { ret; }
    info.size_bits = self.ptr_size * 8;
    info.align_bits = self.ptr_align * 8;
}

// enum 与 struct/union/shape 的布局。字段索引约定（与 ir2vm 一致）：
// enum: 0 = tag，1..n = 各变体载荷（同一偏移）；其余按声明顺序。
fn layout_agg(self: *Layout, is_enum: bool, is_union: bool, agg: Agg, info: *LayoutInfo) {
    var n = agg.variants.len;

    if is_enum {
        if agg.tag_type == 0 { ret; }
        var tag_info = info_at(self, agg.tag_type);
        if tag_info == null || !tag_info.resolved { ret; }

        var tag_align_bits = tag_info.align_bits;
        if tag_align_bits == 0 { tag_align_bits = 8; }
        var payload_align_bits: usize = 8;
        var payload_size_bits: usize = 0;

        var i: usize = 0;
        while i < n {
            var m = member_at(&agg.variants, i);
            if m != null && m.type_id != 0 {
                var fi = info_at(self, m.type_id);
                if fi == null || !fi.resolved { ret; }
                var fa = fi.align_bits;
                if fa == 0 { fa = 8; }
                if fa > payload_align_bits { payload_align_bits = fa; }
                if fi.size_bits > payload_size_bits { payload_size_bits = fi.size_bits; }
            }
            i = i + 1;
        }

        var payload_off = round_up_bits(tag_info.size_bits, payload_align_bits);
        var final_align = tag_align_bits;
        if payload_align_bits > final_align { final_align = payload_align_bits; }
        info.size_bits = round_up_bits(payload_off + payload_size_bits, final_align);
        info.align_bits = final_align;

        push_usize(&info.field_bit_offsets, 0);
        push_usize(&info.field_bit_sizes, tag_info.size_bits);
        i = 0;
        while i < n {
            var m = member_at(&agg.variants, i);
            var fsize: usize = 0;
            if m != null && m.type_id != 0 {
                var fi = info_at(self, m.type_id);
                if fi != null { fsize = fi.size_bits; }
            }
            push_usize(&info.field_bit_offsets, payload_off);
            push_usize(&info.field_bit_sizes, fsize);
            i = i + 1;
        }
        ret;
    }

    if n == 0 {
        info.size_bits = 0;
        info.align_bits = 8;
        ret;
    }

    var max_align_bits: usize = 8;
    var offset_bits: usize = 0;

    var i: usize = 0;
    while i < n {
        var m = member_at(&agg.variants, i);
        var fsize_bits: usize = 0;
        var fsize_align: usize = 8;
        if m != null && m.type_id != 0 {
            var fi = info_at(self, m.type_id);
            if fi == null || !fi.resolved { ret; }
            fsize_bits = fi.size_bits;
            fsize_align = fi.align_bits;
            if fsize_align == 0 { fsize_align = 8; }
        }

        if is_union {
            push_usize(&info.field_bit_offsets, 0);
            push_usize(&info.field_bit_sizes, fsize_bits);
            if fsize_bits > info.size_bits { info.size_bits = fsize_bits; }
        } else {
            offset_bits = round_up_bits(offset_bits, fsize_align);
            push_usize(&info.field_bit_offsets, offset_bits);
            push_usize(&info.field_bit_sizes, fsize_bits);
            offset_bits = offset_bits + fsize_bits;
        }
        if fsize_align > max_align_bits { max_align_bits = fsize_align; }
        i = i + 1;
    }

    if is_union {
        info.size_bits = round_up_bits(info.size_bits, max_align_bits);
    } else {
        info.size_bits = round_up_bits(offset_bits, max_align_bits);
    }
    info.align_bits = max_align_bits;
}

fn resolve_members(self: *Layout, tt: *TypeTable, agg: Agg) {
    var i: usize = 0;
    while i < agg.variants.len {
        var m = member_at(&agg.variants, i);
        if m != null && m.type_id != 0 { resolve_one(self, tt, m.type_id); }
        i = i + 1;
    }
}

fn resolve_one(self: *Layout, tt: *TypeTable, id: usize) {
    var info = info_at(self, id);
    if info == null || info.resolved { ret; }
    var e = node(tt, id);
    if e == null { ret; }

    info.resolved = true;

    match e.node {
        .Error => { info.size_bits = 0; info.align_bits = 8; },
        .Infer => { info.size_bits = 0; info.align_bits = 8; },
        .Forall => { info.size_bits = 0; info.align_bits = 8; },
        .Exists => { info.size_bits = 0; info.align_bits = 8; },
        .Never => { info.size_bits = 0; info.align_bits = 8; },
        .Unit => { info.size_bits = 0; info.align_bits = 8; },
        .Any => { layout_ptr(self, info); },
        .Bool => { layout_scalar(info, 8); },
        .Int(x) => { layout_scalar(info, x.bits); },
        .Float(x) => { layout_scalar(info, x.bits); },
        .Ptr(x) => { layout_ptr(self, info); },
        .Array(a) => {
            resolve_one(self, tt, a.element);
            var elem = info_at(self, a.element);
            if elem == null || !elem.resolved { ret; }
            info.size_bits = elem.size_bits * a.len;
            info.align_bits = elem.align_bits;
        },
        .Struct(agg) => {
            resolve_members(self, tt, agg);
            layout_agg(self, false, false, agg, info);
        },
        .Union(agg) => {
            resolve_members(self, tt, agg);
            layout_agg(self, false, true, agg, info);
        },
        .Enum(agg) => {
            if agg.tag_type != 0 { resolve_one(self, tt, agg.tag_type); }
            resolve_members(self, tt, agg);
            layout_agg(self, true, false, agg, info);
        },
        .Shape(agg) => {
            resolve_members(self, tt, agg);
            layout_agg(self, false, false, agg, info);
        },
        .Fn(f) => {
            var i: usize = 0;
            while i < f.params.len {
                resolve_one(self, tt, usize_at(&f.params, i));
                i = i + 1;
            }
            resolve_one(self, tt, f.ret_type);
            layout_ptr(self, info);
        },
    }
}

fn resolve(self: *Layout, tt: *TypeTable) {
    ensure(self, tt.table.len);
    var i: usize = 1;
    while i < tt.table.len {
        resolve_one(self, tt, i);
        i = i + 1;
    }
}

fn main() {
    var tt: TypeTable = undefined;
    match TypeTable.new() {
        .Ok(v) => { tt = v; },
        .Err(error) => { @assert(false, "layout test: tt init"); ret; },
    }
    defer { TypeTable.drop(&tt); }

    var i32t = make_int(&tt, 32, true);
    var u8t = make_int(&tt, 8, false);

    var members = Vec.blank(@sizeof(Member), @alignof(Member));
    push_member(&members, "a", i32t);
    push_member(&members, "b", u8t);
    var agg = Agg { .mode = .Auto, .fixed_align_bits = 0, .tag_type = 0, .variants = members };
    var sid = make_agg(&tt, .Struct, agg);

    var layout: Layout = undefined;
    match new() {
        .Ok(v) => { layout = v; },
        .Err(error) => { @assert(false, "layout test: layout init"); ret; },
    }
    defer { drop(&layout); }

    resolve(&layout, &tt);
    var info = info_of(&layout, sid);
    if info == null {
        @assert(false, "struct layout missing");
        ret;
    }
    @assert(info.size_bits == 64, "struct size_bits");
    @assert(info.align_bits == 32, "struct align_bits");
    @assert(usize_at(&info.field_bit_offsets, 0) == 0, "field0 off");
    @assert(usize_at(&info.field_bit_offsets, 1) == 32, "field1 off");

    var arr = make_array(&tt, i32t, 3);
    resolve(&layout, &tt);
    var ainfo = info_of(&layout, arr);
    if ainfo == null {
        @assert(false, "array layout missing");
        ret;
    }
    @assert(ainfo.size_bits == 96, "array size_bits");
    @assert(ainfo.align_bits == 32, "array align_bits");

    var pt = TypeTable.make_ptr(&tt, u8t);
    resolve(&layout, &tt);
    var pinfo = info_of(&layout, pt);
    if pinfo == null {
        @assert(false, "ptr layout missing");
        ret;
    }
    @assert(pinfo.size_bits == @sizeof(usize) * 8, "ptr size_bits");

    ret;
}