#include "Python.h"
#include "pyarena.h"

/* An arena list is a linked list that can store PyObjects. */

typedef struct _arena_list {
	struct _arena_list *al_next;
	void *al_pointer;
} PyArenaList;

/* A simple arena block structure */
/* TODO(jhylton): Measurement to justify block size. */

#define DEFAULT_BLOCK_SIZE 8192
typedef struct _block {
	size_t ab_size;
	size_t ab_offset;
	struct _block *ab_next;
	void *ab_mem;
} block;

struct _arena {
	block *a_head;
	block *a_cur;
	PyArenaList *a_object_head;
	PyArenaList *a_object_tail;
};

static PyArenaList*
PyArenaList_New(void) 
{
	PyArenaList *alist = (PyArenaList *)malloc(sizeof(PyArenaList));
	if (!alist)
		return NULL;

	alist->al_next = NULL;
	alist->al_pointer = NULL;
	return alist;
}

static void
PyArenaList_FreeObject(PyArenaList *alist) 
{
	while (alist) {
		PyArenaList *prev;
		Py_XDECREF((PyObject *)alist->al_pointer);
		alist->al_pointer = NULL;
		prev = alist;
		alist = alist->al_next;
		free(prev);
	}
}

static block *
block_new(size_t size)
{
	/* Allocate header and block as one unit. 
	   ab_mem points just past header. */
	block *b = (block *)malloc(sizeof(block) + size);
	if (!b)
		return NULL;
	b->ab_size = size;
	b->ab_mem = (void *)(b + 1);
	b->ab_next = NULL;
	b->ab_offset = 0;
	return b;
}

static void
block_free(block *b) {
	while (b) {
		block *next = b->ab_next;
		free(b);
		b = next;
	}
}

static void *
block_alloc(block *b, size_t size)
{
	void *p;
	assert(b);
	if (b->ab_offset + size > b->ab_size) {
		/* If we need to allocate more memory than will fit in
		   the default block, allocate a one-off block that is
		   exactly the right size. */
		/* TODO(jhylton): Think about space waste at end of block */
		block *new = block_new(
				size < DEFAULT_BLOCK_SIZE ?
				DEFAULT_BLOCK_SIZE : size);
		if (!new)
			return NULL;
		assert(!b->ab_next);
		b->ab_next = new;
		b = new;
	}

	assert(b->ab_offset + size <= b->ab_size);
	p = (void *)(((char *)b->ab_mem) + b->ab_offset);
	b->ab_offset += size;
	return p;
}

PyArena *
PyArena_New()
{
	PyArena* arena = (PyArena *)malloc(sizeof(PyArena));
	if (!arena)
		return NULL;

	arena->a_head = block_new(DEFAULT_BLOCK_SIZE);
	arena->a_cur = arena->a_head;
	arena->a_object_head = PyArenaList_New();
	arena->a_object_tail = arena->a_object_head;
	return arena;
}

void
PyArena_Free(PyArena *arena)
{
	assert(arena);
	block_free(arena->a_head);
	PyArenaList_FreeObject(arena->a_object_head);
	free(arena);
}

void *
PyArena_Malloc(PyArena *arena, size_t size) 
{
	void *p = block_alloc(arena->a_cur, size);
	if (!p)
		return NULL;
	/* Reset cur if we allocated a new block. */
	if (arena->a_cur->ab_next) {
		arena->a_cur = arena->a_cur->ab_next;
	}
	return p;
}

int
PyArena_AddPyObject(PyArena *arena, PyObject *pointer) 
{
	PyArenaList *tail = arena->a_object_tail;
	assert(pointer);
	tail->al_next = PyArenaList_New();
	tail->al_pointer = pointer;
	arena->a_object_tail = tail->al_next;
	return 1;
}