#include <stdint.h>
#include <stdlib.h>
#include <stddef.h>
#include <stdio.h>
#include <string.h>
#include <sys/mman.h>
#include <errno.h>

#define BLOCK_SIZE 16

typedef struct Header Header;
typedef struct Arena Arena;

extern char _end[];

/* Structure for a Memory Header */
struct Header {
	Header *prev, *next;
	size_t size;
	char magic[4];
};

/* Structure for a Memory Arena */
struct Arena {
	Arena *next;
	void *start, *end;
	size_t size;
};

Arena primaryArena = {
	.next = NULL,
	.start = NULL,
	.end = NULL,
	.size = 0
};

/* Allocate a region of memory in an arena */
static void *
arena_malloc(Arena *arena, size_t size)
{
	size_t blockSize, gapSize;
	intptr_t blockEnd;
	Header *prev, *head, *next;
	head = prev = arena->start;
	next = NULL;

	/* Minimum allocation */
	if (size % BLOCK_SIZE)
		size += BLOCK_SIZE - (size % BLOCK_SIZE);

	/* Block does not exist, create heap */
	if (head->prev != head || memcmp(head->magic, "HEAP", 4)) {
		head->prev = head;
		head->next = NULL;
		head->size = size;
		memcpy(head->magic, "HEAP", 4);
		memset((void *) (head + 1), 0, size);
		return (void *) (head + 1);
	}

	/* Find gap */
	while (head->next) {
		next = head->next;
		blockSize = sizeof(Header) + head->size;
		blockEnd = (uintptr_t) head + blockSize;
		gapSize = (size_t) next - blockEnd;
		prev = head;

		/* Fit in gap */
		if (gapSize >= size + sizeof(Header)) {
			head = (void *) blockEnd;
			head->next = next;
			head->prev = prev;
			prev->next = head;
			next->prev = head;
			head->size = size;
			memcpy(head->magic, "HEAP", 4);
			memset((void *) (head + 1), 0, size);
			return (void *) (head + 1);
		}

		head = head->next;
	}

	/* Add to end */
	blockSize = sizeof(Header) + head->size;
	blockEnd = (uintptr_t) head + blockSize;
	/* Ensure more arenas can be allocated */
	gapSize = (size_t) arena->end - blockEnd - sizeof(Arena);
	/* Fit in gap */
	if (gapSize >= size + sizeof(Header)) {
		prev = head;
		head = (void *) blockEnd;
		head->next = NULL;
		head->prev = prev;
		prev->next = head;
		head->size = size;
		memcpy(head->magic, "HEAP", 4);
		memset((void *) (head + 1), 0, size);
		return (void *) (head + 1);
	}

	return NULL;
}

/* Allocate a region of memory */
void *
malloc(size_t size)
{
	void *addr = NULL;
	Arena *arena;
	size_t arenaSize = 64 * 1024;
	while (arenaSize < size + sizeof(Arena))
		arenaSize += 64 * 1024;

	for (arena = &primaryArena; !addr; arena = arena->next) {
		/* Allocate a new arena */
		if (!arena->size) {
			arena->size = arenaSize;
			arena->start = mmap(NULL, arena->size,
			                    PROT_READ | PROT_WRITE,
			                    MAP_PRIVATE | MAP_ANONYMOUS, 0, 0);
			if (arena->start == MAP_FAILED) {
				errno = ENOMEM;
				return NULL;
			}
			arena->end = arena->start + arena->size;
		}

		/* Attempt to use arena */
		addr = arena_malloc(arena, size);
		if (!addr && !arena->next)
			arena->next = arena->start + arena->size
			            - sizeof(Arena);
	}

	return addr;
}

/* Free an allocated region of memory */
void
free(void *addr)
{
	Header *prev, *head, *next;
	head = (Header *) addr - 1;
	if (memcmp(head->magic, "HEAP", 4)) {
		printf("free(): invalid pointer\n");
		abort();
	}
	prev = head->prev;
	next = head->next;
	memset(head, 0, sizeof(Header));

	if (prev != head)
		prev->next = next;
	if (next)
		next->prev = prev;
}

/* Re-allocate a region of memory */
void *
realloc(void *addr, size_t size)
{
	void *ptr = malloc(size);
	if (addr) {
		Header *old = (Header *) addr - 1;
		if (memcmp(old->magic, "HEAP", 4)) {
			printf("realloc(): invalid pointer\n");
			abort();
		}
		memcpy(ptr, addr, old->size);
		free(addr);
	}
	return ptr;
}