#include "internal.h"

#define ARENA_ALIGN 8
#define ARENA_BLOCK_SIZE 1024 * 1024

#define ALIGN_DOWN(n, a) ((n) & ~((a)-1))
#define ALIGN_UP(n, a) ALIGN_DOWN((n) + (a)-1, (a))
#define ALIGN_DOWN_PTR(p, a) ((void*)ALIGN_DOWN((uintptr)(p), (a)))
#define ALIGN_UP_PTR(p, a) ((void*)ALIGN_UP((uintptr)(p), (a)))

struct Block
{
    struct Block *next;
    byte   buf[];
};

struct mem·Arena
{
    void  *heap;
    mem·Allocator mem;

    byte  *off;
    byte  *end;
    struct Block *curr;
    struct Block first;
};

static
void*
·arenaalloc(void *heap, uint n, ulong size)
{
    return mem·arenaalloc(heap, n, size);
}

static
void
·arenafree(void *heap, void *ptr)
{
    /* no-op */
}

mem·Allocator mem·ArenaAllocator = {
    .alloc = ·arenaalloc,
    .free  = ·arenafree,
};


static
void
grow(mem·Arena *a, vlong min)
{
    uintptr size;
    struct Block *blk;

    size   = ALIGN_UP(MAX(min, ARENA_BLOCK_SIZE), ARENA_ALIGN);
    blk    = a->mem.alloc(a->heap, 1, sizeof(*blk) + size);
    a->off = blk->buf;
    a->end = a->off + size;

    assert(a->curr->next == nil);
    assert(a->off == ALIGN_DOWN_PTR(a->off, ARENA_ALIGN));

    a->curr->next = blk;
    a->curr       = blk;
}

mem·Arena*
mem·makearena(mem·Allocator from, void *impl)
{
    mem·Arena *a  = from.alloc(impl, 1, sizeof(*a) + ARENA_BLOCK_SIZE);
    a->mem        = from;
    a->heap       = impl;
    a->off        = a->first.buf;
    a->end        = a->first.buf + ARENA_BLOCK_SIZE;
    a->curr       = &a->first;
    a->first.next = nil;

    return a;
}

void
mem·freearena(mem·Arena *a)
{
    struct Block *it, *next;

    it = a->first.next;
    while (it != nil) {
        next = it->next;
        a->mem.free(a->heap, it);
        it   = next;
    }

    a->mem.free(a->heap, a);
}

void*
mem·arenaalloc(mem·Arena *a, uint n, ulong size)
{
    if(!n) {
        return nil;
    }

    void *ptr;
    // TODO(nnoll): check for overflow
    size = n * size;

    if (size > (ulong)(a->end - a->off)) {
        grow(a, size);
        assert(size <= (uintptr)(a->end - a->off));
    }

    ptr    = a->off;
    a->off = ALIGN_UP_PTR(a->off + size, ARENA_ALIGN);

    assert(a->off <= a->end);
    assert(ptr == ALIGN_DOWN_PTR(ptr, ARENA_ALIGN));

    return ptr;
}