#import "std.sx";

// --- Allocator protocol ---

Allocator :: struct {
    ctx: *void;
    alloc_fn: (*void, s64) -> *void;
    free_fn: (*void, *void) -> void;
}

allocator_alloc :: (a: Allocator, size: s64) -> *void {
    a.alloc_fn(a.ctx, size);
}

allocator_dealloc :: (a: Allocator, ptr: *void) {
    a.free_fn(a.ctx, ptr);
}

alloc :: ufcs allocator_alloc;
dealloc :: ufcs allocator_dealloc;

// --- GPA: general purpose allocator (malloc/free wrapper) ---

GPA :: struct {
    alloc_count: s64;
}

gpa_alloc :: (ctx: *void, size: s64) -> *void {
    gpa : *GPA = xx ctx;
    gpa.alloc_count += 1;
    malloc(size);
}

gpa_free :: (ctx: *void, ptr: *void) {
    gpa : *GPA = xx ctx;
    gpa.alloc_count -= 1;
    free(ptr);
}

gpa_create :: (gpa: *GPA) -> Allocator {
    Allocator.{ ctx = gpa, alloc_fn = gpa_alloc, free_fn = gpa_free };
}

// --- Arena: multi-chunk bump allocator (Zig-inspired) ---

ArenaChunk :: struct {
    next: *ArenaChunk;
    cap: s64;
}

Arena :: struct {
    first: *ArenaChunk;
    end_index: s64;
    parent: Allocator;
}

arena_add_chunk :: (a: *Arena, min_size: s64) {
    prev_cap := if a.first != null then a.first.cap else 0;
    needed := min_size + 16 + 16;
    len := (prev_cap + needed) * 3 / 2;
    raw := a.parent.alloc(len);
    chunk : *ArenaChunk = xx raw;
    chunk.next = a.first;
    chunk.cap = len;
    a.first = chunk;
    a.end_index = 0;
}

arena_alloc :: (ctx: *void, size: s64) -> *void {
    a : *Arena = xx ctx;
    aligned := (size + 7) & (0 - 8);
    if a.first != null {
        usable := a.first.cap - 16;
        if a.end_index + aligned <= usable {
            buf : [*]u8 = xx a.first;
            ptr := @buf[16 + a.end_index];
            a.end_index = a.end_index + aligned;
            return ptr;
        }
    }
    arena_add_chunk(a, aligned);
    buf : [*]u8 = xx a.first;
    ptr := @buf[16 + a.end_index];
    a.end_index = a.end_index + aligned;
    ptr;
}

arena_free :: (ctx: *void, ptr: *void) {
}

arena_create :: (a: *Arena, parent: Allocator, size: s64) -> Allocator {
    a.first = null;
    a.end_index = 0;
    a.parent = parent;
    arena_add_chunk(a, size);
    Allocator.{ ctx = a, alloc_fn = arena_alloc, free_fn = arena_free };
}

arena_reset :: (a: *Arena) {
    // Keep first chunk (newest/largest), free the rest
    if a.first != null {
        it := a.first.next;
        while it != null {
            next := it.next;
            a.parent.dealloc(it);
            it = next;
        }
        a.first.next = null;
    }
    a.end_index = 0;
}

arena_deinit :: (a: *Arena) {
    it := a.first;
    while it != null {
        next := it.next;
        a.parent.dealloc(it);
        it = next;
    }
    a.first = null;
    a.end_index = 0;
}

// --- BufAlloc: bump allocator backed by a user-provided slice ---

BufAlloc :: struct {
    buf: [*]u8;
    len: s64;
    pos: s64;
}

buf_alloc :: (ctx: *void, size: s64) -> *void {
    b : *BufAlloc = xx ctx;
    aligned := (size + 7) & (0 - 8);
    if b.pos + aligned > b.len {
        return null;
    }
    ptr := @b.buf[b.pos];
    b.pos = b.pos + aligned;
    ptr;
}

buf_free :: (ctx: *void, ptr: *void) {
}

buf_create :: (b: *BufAlloc, buf: [*]u8, len: s64) -> Allocator {
    b.buf = buf;
    b.len = len;
    b.pos = 0;
    Allocator.{ ctx = b, alloc_fn = buf_alloc, free_fn = buf_free };
}

buf_reset :: (b: *BufAlloc) {
    b.pos = 0;
}