// VERSION: 3.0.0
/* This file provides an implementation of an arena allocator that allows for
 * many allocations, which are then freed all at once.  It is particularly
 * useful for functions with lots of temporary allocations (string parsing and
 * modification?) as you don't have to keep track of every little allocation,
 * and can instead free all memory used afterwards.
 *
 * This implementation is very loosely based off of Zig's arena allocator:
 * https://github.com/ziglang/zig/blob/master/lib/std/heap/arena_allocator.zig
 *
 *
 * Copyright (c) 2021 nytpu <alex [at] nytpu.com>
 * SPDX-License-Identifier: MPL-2.0
 * The orginal source for this file is available at <https://git.sr.ht/~nytpu/arena.c>.
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at <https://mozilla.org/MPL/2.0/>.
 */

#include "arena.h"
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>

// An arena "page" is just a buffer of size cap with sz amount of data
// contained in it
struct page {
	size_t sz; // size of the buf's contained data, not the full buf size
	size_t cap;
	uint8_t *buf;
};

// An arena is just a set of buffer pages (len pages long).  The arenas are
// usually of size sys_pgsz, but may be larger.
struct arena {
	size_t sys_pgsz;
	void *(*calloc)(size_t, size_t);
	void *(*realloc)(void *, size_t);
	void (*free)(void *);

	size_t len;
	struct page *pg; // array of pages
};


// I usually avoid using these system detect macros and for this case it's
// probably not worth it, but I'm gonna do it anyways
// TODO: figure out how to get page size on more systems
#if defined(_POSIX_C_SOURCE) && _POSIX_C_SOURCE >= 200112L
#include <unistd.h>
static size_t
get_page_size(void)
{
	return (size_t)sysconf(_SC_PAGESIZE);
}
#elif defined(_BSD_SOURCE)
#include <unistd.h>
static size_t
get_page_size(void)
{
	return (size_t)getpagesize();
}
#elif defined(_WIN32)
#include <windows.h>
static size_t
get_page_size(void)
{
	SYSTEM_INFO si;
	GetSystemInfo(&si);
	return (size_t)si.dwPageSize;
}
#else
static size_t
get_page_size(void)
{
	return ARENA_DEFAULT_PAGE_SIZE;
}
#endif

// Add a new, empty page of size to an arena.  Returns 0 if successful, -1 on
// failure
static int
arena_new_page(arena *a, size_t size)
{
	// do this runaround to prevent nullifying the valid a->pg on
	// allocation failure
	struct page *p = a->realloc(a->pg, (a->len + 1) * sizeof(*a->pg));
	if (!p) return -1;
	a->pg = p;

	a->pg[a->len].sz = 0;
	a->pg[a->len].cap = size;
	a->pg[a->len].buf = a->calloc(size, sizeof(*a->pg[a->len].buf));
	if (!a->pg[a->len].buf) return -1;

	a->len++;
	return 0;
}

arena *
arena_init_with_allocator(
	void *(*_calloc)(size_t, size_t),
	void *(*_realloc)(void *, size_t),
	void (*_free)(void *)
) {
	if (_calloc == NULL) _calloc = calloc;
	arena *a = _calloc(1, sizeof(*a));
	if (!a) return NULL;
	a->sys_pgsz = get_page_size();

	a->calloc = _calloc;
	if (_realloc != NULL) a->realloc = _realloc;
	else a->realloc = realloc;
	if (_free != NULL) a->free = _free;
	else a->free = free;

	a->len = 0;
	a->pg = NULL;
	int rc = arena_new_page(a, a->sys_pgsz);
	if (rc == -1) {
		arena_deinit(&a);
		return NULL;
	}
	return a;
}

arena *
arena_init(void)
{
	return arena_init_with_allocator(calloc, realloc, free);
}

void
arena_deinit(arena **a)
{
	if (a == NULL || *a == NULL) return;

	void (*_free)(void *) = (*a)->free;
	for (size_t i = 0; i < (*a)->len; ++i) _free((*a)->pg[i].buf);
	_free((*a)->pg);
	_free(*a);
	*a = NULL;
}

void
arena_clear(arena *a)
{
	if (a == NULL) return;
	for (size_t i = 0; i < a->len; ++i) a->pg[i].sz = 0;
}

void *
arena_alloc(arena *a, size_t size)
{
	if (a == NULL) return NULL;
	for (size_t i = 0; i < a->len; ++i) {
		if ((a->pg[i].cap - a->pg[i].sz) >= size) {
			a->pg[i].sz += size;
			return a->pg[i].buf + (a->pg[i].sz - size);
		}
	}

	size_t newsz = (size > a->sys_pgsz) ? size : a->sys_pgsz;
	int rc = arena_new_page(a, newsz);
	if (rc == -1) return NULL;
	a->pg[a->len - 1].sz += newsz;
	return a->pg[a->len - 1].buf;
}


#ifdef TEST
#include "testing.h"

int
main(void)
{
	begin_tests("arena.c");

	begin_group("arena_init");
	arena *a = arena_init();
	assert(a != NULL);
	printf("\tDetected page size: %ld\n", a->sys_pgsz);
	// checking for sensible defaults, not really necessary
	assert(a->len == 1);
	assert(a->pg != NULL);
	assert(a->pg[0].sz == 0);
	assert(a->pg[0].cap == a->sys_pgsz);
	assert(a->pg[0].buf != NULL);
	end_group();

	begin_group("arena_alloc");
	void *dummy = arena_alloc(a, 10);
	assert(dummy == a->pg[0].buf);
	assert(a->pg[0].sz == 10);

	dummy = arena_alloc(a, a->sys_pgsz + 1);
	assert(a->len == 2);
	assert(a->pg[1].sz == a->sys_pgsz + 1);
	assert(a->pg[1].cap == a->sys_pgsz + 1);
	assert(dummy == a->pg[1].buf);

	dummy = arena_alloc(a, 20);
	assert(dummy == (a->pg[0].buf + 10));
	assert(a->pg[0].sz == 30);
	end_group();

	begin_group("arena_clear");
	arena_clear(a);
	assert(a->len == 2);
	assert(a->pg[0].sz == 0);
	assert(a->pg[1].sz == 0);
	end_group();

	arena_deinit(&a);
	end_tests();
	return 0;
}

#endif // TEST