// Expect retcode = 0
// 分块 bump 分配器：按块分配、块内顺序推进，只在 drop 时整体释放。
// 适合 AST/IR 这类生命周期一致的批量小对象；不支持单独 free。
// 失败（内存不足/参数非法）返回 null，不改变已有分配。

type .{ .Vec } = @import("vec.spl");
type Mem = @import("../sys/mem.spl");

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

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 is_pow2(x: usize): bool {
    ret x != 0 && (x & (x - 1)) == 0;
}

fn align_up(x: usize, align: usize): usize {
    var mask = align - 1;
    ret (x + mask) & ~mask;
}

fn chunk_at(self: *Arena, index: usize): *Chunk {
    match Vec.get(&self.chunks, index) {
        .Ok(ptr) => { ret ptr as *Chunk; },
        .Err(error) => { ret null; },
    }
    ret null;
}

fn push_chunk(self: *Arena, size: usize): bool {
    var base = Mem.alloc(size) as *u8;
    if base == null { ret false; }
    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 false; },
    }
    self.current = self.chunks.len - 1;
    ret true;
}

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

    while true {
        if self.current >= self.chunks.len {
            var want = self.chunk_size;
            var need = size + align;
            if need > want { want = need; }
            if !push_chunk(self, want) { ret null; }
        }

        var c = chunk_at(self, self.current);
        if c == null { ret null; }
        var addr = c.base as usize;
        var aligned = align_up(addr + c.used, align);
        var offset = aligned - addr;
        if offset <= c.size && size <= c.size - offset {
            c.used = offset + size;
            ret aligned as *u8;
        }
        self.current = self.current + 1;
    }
    ret null;
}

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 test_alloc_and_alignment() {
    var a = new(64);
    @assert(chunk_count(a) == 0, "lazy");

    var p1 = alloc(&a, 1, 1);
    var p8 = alloc(&a, 16, 8);
    var p16 = alloc(&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(&a, 1, 1);
    @assert(q == p1, "reset reuses first chunk");

    drop(&a);
    @assert(chunk_count(a) == 0, "drop empties");
}

fn test_growth() {
    var a = new(64);
    var i: i32 = 0;
    while i < 100 {
        var p = alloc(&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(&a, 1024, 16);
    @assert(big != null && is_aligned(big, 16), "oversized alloc");
    drop(&a);
}

fn test_failures() {
    var a = new(64);
    @assert(alloc(&a, 0, 1) == null, "zero size");
    var d = alloc(&a, 8, 0);
    @assert(d != null && is_aligned(d, 1), "align 0 defaults to 1");
    @assert(alloc(&a, 8, 3) == null, "non-pow2 align");
    drop(&a);
}

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