#include <arena.h>
#include <string.h>
#include <os_align.h>
#include "os_macros.h"


#define IS_POWER_OF_2(n)        ((n) != 0 && (((n) & ((n) - 1)) == 0))
#define DEFAULT_ALIGNMENT       OS_ALIGN_SIZE

/* ------------------------------------------------------------------------------------------------------------------ */
/*  */


static
os_uintptr_t align_forward(os_uintptr_t ptr, os_size_t align) {
    os_uintptr_t p, a, modulo;

    OS_ASSERT(IS_POWER_OF_2(align));

    p = ptr;
    a = (os_uintptr_t )align;
    // Same as (p % a) but faster as 'a' is a power of two
    modulo = p & (a-1);

    if (modulo != 0) {
        // If 'p' address is not aligned, push the address to the
        // next value which is aligned
        p += a - modulo;
    }
    return p;
}


static void *arena_alloc_align(arena_t *a, os_size_t size, os_size_t align) {
    // Align 'curr_offset' forward to the specified alignment
    os_uintptr_t curr_ptr = (os_uintptr_t) a->buf + (os_uintptr_t) a->curr_offset;
    os_uintptr_t offset = align_forward(curr_ptr, align);
    offset -= (os_uintptr_t) a->buf; // Change to relative offset

    // Check to see if the backing memory has space left
    if (offset+size <= a->buf_size) {
        void *ptr = &a->buf[offset];
        a->prev_offset = offset;
        a->curr_offset = offset+size;

        // Zero new memory by default
        memset(ptr, 0, size);
        return ptr;
    }
    // Return NULL if the arena is out of memory (or handle differently)
    return NULL;
}

static void *arena_resize_align(arena_t *a, void *old_memory, size_t old_size, size_t new_size, size_t align) {
    unsigned char *old_mem = (unsigned char *)old_memory;

    OS_ASSERT(IS_POWER_OF_2(align));

    if (old_mem == NULL || old_size == 0) {
        return arena_alloc_align(a, new_size, align);
    } else if (a->buf <= old_mem && old_mem < a->buf+a->buf_size) {
        if (a->buf+a->prev_offset == old_mem) {
            a->curr_offset = a->prev_offset + new_size;
            if (new_size > old_size) {
                // Zero the new memory by default
                memset(&a->buf[a->curr_offset], 0, new_size-old_size);
            }
            return old_memory;
        } else {
            void *new_memory = arena_alloc_align(a, new_size, align);
            size_t copy_size = old_size < new_size ? old_size : new_size;
            // Copy across old memory to the new memory
            memmove(new_memory, old_memory, copy_size);
            return new_memory;
        }

    } else {
        OS_ASSERT(0 && "Memory is out of bounds of the buffer in this arena");
        return NULL;
    }

}

/* ------------------------------------------------------------------------------------------------------------------ */
/*  */

void arena_init(arena_t* self, uint8_t * buf, os_size_t buf_size){
    self->buf = buf;
    self->buf_size = buf_size;
    self->prev_offset = 0;
    self->curr_offset = 0;
}

void* arena_alloc(arena_t* self, os_size_t size){
    return arena_alloc_align(self, size, DEFAULT_ALIGNMENT);
}

void* arena_realloc(arena_t* self, void* old, os_size_t old_size, os_size_t new_size){
    return arena_resize_align(self, old, old_size, new_size, DEFAULT_ALIGNMENT);
}

void arena_free(arena_t* self, void* ptr)
{}

void arean_destroy(arena_t* self){
    self->curr_offset = 0;
    self->prev_offset = 0;
}

/* ------------------------------------------------------------------------------------------------------------------ */
/*  */


temp_arena_memory_t temp_arena_memory_begin(arena_t *a) {
    temp_arena_memory_t temp;
    temp.arena = a;
    temp.prev_offset = a->prev_offset;
    temp.curr_offset = a->curr_offset;
    return temp;
}

void temp_arena_memory_end(temp_arena_memory_t temp) {
    temp.arena->prev_offset = temp.prev_offset;
    temp.arena->curr_offset = temp.curr_offset;
}