#include <cstdio>
#include <cstring>
#include <cassert>
#include <cstdlib>

#include "memory-arena.hpp"

static bool is_power_of_two(uintptr_t x) {
	return (x & (x-1)) == 0;
}

static uintptr_t align_forward(uintptr_t ptr, size_t align) {
	assert(is_power_of_two(align));

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

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

void arena_init(MemoryArena *arena, size_t size) {
	arena->buffer = (uint8_t*)malloc(size);
	arena->size = size;
	arena->offset = 0;
}

void arena_free(MemoryArena *arena) {
	free(arena->buffer);
	arena->buffer = NULL;
	arena->size = 0;
}

void* arena_malloc(MemoryArena *arena, size_t size, size_t align) {
	uintptr_t curr_ptr = (uintptr_t)arena->buffer + arena->offset;
	uintptr_t offset = align_forward(curr_ptr, align) - (uintptr_t)arena->buffer;

	if (offset + size >= arena->size) {
		// TODO: grow arena
		assert(false && "Arena ran out of space");
		return NULL;
	}

	uint8_t *ptr = &arena->buffer[offset];
	arena->offset = offset + size;

	return ptr;
}

void* arena_calloc(MemoryArena *arena, size_t size, size_t align) {
	void* ptr = arena_malloc(arena, size, align);
	if (ptr != NULL) {
		memset(ptr, 0, size);
	}
	return ptr;
}

void arena_clear(MemoryArena *arena) {
	arena->offset = 0;
}