// Expect retcode = 0

// 分块 bump 分配器：按块分配、块内顺序推进，只在 drop 时整体释放。
// 适合 AST/IR 这类生命周期一致的批量小对象；不支持单独 free。
// 分配失败通过 ArenaResult 区分原因：size==0/非法对齐/内存不足/地址溢出。
// 失败的 alloc 不改变已有分配。

type .{ .Vec } = @import("vec.spl");
type .{ .Error, .Kind, .Status } = @import("error.spl");
type .{ .checked_add, .align_up } = @import("memory.spl");
type Mem = @import("../sys/mem.spl");

type Chunk = struct {
    base: *u8,
    size: usize,
    used: usize,
}

type ArenaResult = enum {
    Ok: *u8,
    Err: Error,
}

type Arena = @self;

chunk_size: usize,
chunks: Vec, // 元素类型 Chunk
current: usize,

fn new(chunk_size: usize): Arena {
    var cs = chunk_size;
    if cs == 0 { cs = 4096; }
    ret Arena {
        .chunk_size = cs,
        .chunks = Vec.blank(@sizeof(Chunk), 8),
        .current = 0,
    };
}

fn chunk_at(self: *Arena, index: usize): *Chunk {
    match Vec.get(&self.chunks, index) {
        .Ok(ptr) => { ret ptr as *Chunk; },
        // 越界/无效 Arena 统一返回 null；调用方负责判空，无资源需要释放。
        .Err(error) => { ret null; },
    }
    ret null;
}

fn push_chunk(self: *Arena, size: usize): Status {
    var base = Mem.alloc(size) as *u8;
    if base == null { ret .Err(Error.make(.Overflow)); }
    var c = Chunk { .base = base, .size = size, .used = 0 };
    match Vec.push(&self.chunks, &c as *u8) {
        .Ok => {},
        .Err(error) => { Mem.free(base as *any); ret .Err(error); },
    }
    self.current = self.chunks.len - 1;
    ret .Ok;
}

// 分配 size 字节，地址按 align（非零二次幂）对齐。失败返回 ArenaResult.Err。
fn alloc(self: *Arena, size: usize, align: usize): ArenaResult {
    if size == 0 { ret .Err(Error.make(.InvalidArgument)); }
    if align == 0 { align = 1; }

    while true {
        if self.current >= self.chunks.len {
            var want = self.chunk_size;
            var need: usize = 0;
            match checked_add(size, align) {
                .Ok(value) => { need = value; },
                .Err(error) => { ret .Err(error); },
            }
            if need > want { want = need; }
            match push_chunk(self, want) {
                .Ok => {},
                .Err(error) => { ret .Err(error); },
            }
        }

        var c = chunk_at(self, self.current);
        if c == null { ret .Err(Error.make(.InvalidArgument)); }
        var addr = c.base as usize;
        var used_addr: usize = 0;
        match checked_add(addr, c.used) {
            .Ok(value) => { used_addr = value; },
            .Err(error) => { ret .Err(error); },
        }
        var aligned: usize = 0;
        match align_up(used_addr, align) {
            .Ok(value) => { aligned = value; },
            .Err(error) => { ret .Err(error); }, // 非二次幂对齐或地址溢出
        }
        var offset = aligned - addr;
        if offset <= c.size && size <= c.size - offset {
            c.used = offset + size;
            ret .Ok(aligned as *u8);
        }
        self.current = self.current + 1;
    }
    ret .Err(Error.make(.OutOfRange)); // 不可达：循环内所有路径均已返回
}

fn chunk_count(self: Arena): usize {
    ret self.chunks.len;
}

// 复用已分配的块，不释放内存。
fn reset(self: *Arena) {
    var i: usize = 0;
    while i < self.chunks.len {
        var c = chunk_at(self, i);
        if c != null { c.used = 0; }
        i = i + 1;
    }
    self.current = 0;
}

fn drop(self: *Arena) {
    var i: usize = 0;
    while i < self.chunks.len {
        var c = chunk_at(self, i);
        if c != null { Mem.free(c.base as *any); }
        i = i + 1;
    }
    Vec.drop(&self.chunks);
    self.current = 0;
}

// ---- 测试辅助 ----

fn is_aligned(p: *u8, align: usize): bool {
    ret ((p as usize) & (align - 1)) == 0;
}

fn alloc_ok(a: *Arena, size: usize, align: usize): *u8 {
    match alloc(a, size, align) {
        .Ok(ptr) => { ret ptr; },
        .Err(error) => { @assert(false, "unexpected arena alloc error"); },
    }
    ret null;
}

fn expect_alloc_error(result: ArenaResult, expected: Kind) {
    match result {
        .Ok(ptr) => { @assert(false, "expected arena alloc error"); },
        .Err(error) => {
            var want: Kind = expected;
            @assert(@enum2int(error.kind) == @enum2int(want), "arena alloc error kind");
        },
    }
}

fn test_alloc_and_alignment() {
    var a = new(64);
    @assert(chunk_count(a) == 0, "lazy");

    var p1 = alloc_ok(&a, 1, 1);
    var p8 = alloc_ok(&a, 16, 8);
    var p16 = alloc_ok(&a, 1, 16);
    @assert(p1 != null && p8 != null && p16 != null, "alloc");
    @assert(p1 != p8 && p8 != p16 && p1 != p16, "distinct");
    @assert(is_aligned(p8, 8), "align 8");
    @assert(is_aligned(p16, 16), "align 16");

    p1[0] = 7;
    @assert(p1[0] == 7, "writable");
    @assert(p8[0] == p8[0], "no crash");

    @assert(chunk_count(a) == 1, "one chunk");
    reset(&a);
    @assert(chunk_count(a) == 1, "reset keeps chunks");
    var q = alloc_ok(&a, 1, 1);
    @assert(q == p1, "reset reuses first chunk");

    // drop 后还有断言，故保持手动释放，不改为 defer。
    drop(&a);
    @assert(chunk_count(a) == 0, "drop empties");
}

fn test_growth() {
    var a = new(64);
    defer { drop(&a); }
    var i: i32 = 0;
    while i < 100 {
        var p = alloc_ok(&a, 16, 8);
        @assert(p != null, "grow alloc");
        @assert(is_aligned(p, 8), "grow align");
        p[0] = i as u8;
        i = i + 1;
    }
    @assert(chunk_count(a) > 1, "grew multiple chunks");

    // 超过块大小的单次分配也能成功（新块按需放大）。
    var big = alloc_ok(&a, 1024, 16);
    @assert(big != null && is_aligned(big, 16), "oversized alloc");
}

fn test_failures() {
    var a = new(64);
    defer { drop(&a); }
    expect_alloc_error(alloc(&a, 0, 1), .InvalidArgument);
    var d = alloc_ok(&a, 8, 0);
    @assert(d != null && is_aligned(d, 1), "align 0 defaults to 1");
    expect_alloc_error(alloc(&a, 8, 3), .InvalidAlignment);
}

fn main() {
    test_alloc_and_alignment();
    test_growth();
    test_failures();
}