#include "BumpAlloc.h"
#include "Types.h"
#include <stdio.h>
#include <stdlib.h>

struct BumpAlloc {
    void* memory;
    usize size;
    usize top;
};

static inline usize*
top_size(struct BumpAlloc* b) {
    return (usize*)(b->memory + b->top);
}

static inline int
bump_resizeable(struct BumpAlloc *b, void *ptr) {
    if (NULL == b)
        return 0;

    return b->memory + b->top - *top_size(b) <= ptr &&
           ptr < b->memory + b->top;
}

Result
bump_free(MemoryManager mm, void **ptr) {
    if (NULL == mm.priv)
        return R_Uninitialized;

    struct BumpAlloc *b = (struct BumpAlloc*)mm.priv;

    if (NULL != ptr) {
        if (bump_resizeable(b, *ptr))
            *ptr = NULL;
        else
            return R_InvalidArgument;
    }

    if (b->top > 0) {
        b->top -= *top_size(b);
        if (b->top > 0)
            b->top -= sizeof(usize);
    }

    return R_Ok;
}

Result
bump_destroy(MemoryManager *mm) {
    if (NULL == mm)
        return R_NullArgument;

    if (NULL == mm->priv)
        return R_Uninitialized;

    struct BumpAlloc b = *(struct BumpAlloc*)mm->priv;

    if (NULL == b.memory) {
        free(mm->priv);
        mm->priv = NULL;
        return R_InvalidArgument;
    }

    free(b.memory);
    free(mm->priv);
    mm->priv = NULL;

    return R_Ok;
}

Result
bump_alloc(MemoryManager mm, void **ptr, usize size) {
    if (0 == size || NULL == ptr)
        return R_NullArgument;

    if (NULL == mm.priv)
        return R_Uninitialized;

    struct BumpAlloc *b = (struct BumpAlloc*)mm.priv;

    if (size + b->top + sizeof(usize) > b->size)
        return R_OutOfMemory;

    if (b->top > 0)
        b->top += sizeof(usize);
    usize start = b->top;

    b->top += size;
    // Keep allocated memory size at the end of every allocation
    *top_size(b) = size;

    *ptr = b->memory + start;
    return R_Ok;
}

Result
bump_resize(MemoryManager mm, void **ptr, usize size) {
    if (NULL == mm.priv)
        return R_Uninitialized;

    if (size == 0)
        return bump_free(mm, ptr);

    if (NULL == ptr)
        return R_NullArgument;

    struct BumpAlloc *b = (struct BumpAlloc*)mm.priv;

    if (NULL == ptr || !bump_resizeable(b, *ptr))
        return R_InvalidArgument;

    usize last_size = *top_size(b);
    usize start = b->top - last_size;

    if (start + size + sizeof(usize) > b->size)
        return R_OutOfMemory;

    *top_size(b) = 0;
    b->top = start + size;
    *top_size(b) = size;

    return R_Ok;
}

MemoryManager
M_bump(usize size) {
    struct BumpAlloc *b = malloc(sizeof(struct BumpAlloc));
    b->memory = malloc(size);
    b->size = size;
    b->top = 0;

    return (MemoryManager){
        .alloc = &bump_alloc,
        .realloc = &bump_resize,
        .free = &bump_free,
        .destroy = &bump_destroy,
        .priv = b,
    };
}
