#ifndef ZONE_H
#define ZONE_H

#include <stdint.h>
#include <stddef.h>

#ifndef ZONE_MIN_CAPACITY
#define ZONE_MIN_CAPACITY 1024
#endif

#define ZONE_USE_MALLOC 0
#define ZONE_USE_MMAP   1

typedef struct ZoneFrame ZoneFrame;
struct ZoneFrame {
	ZoneFrame *prev, *next;
	uint8_t *beg, *end;
	uint8_t data[];
};

typedef struct {
	ZoneFrame *cur, *tail;
} Zone;

typedef struct {
	ZoneFrame *zf;
	uint8_t *beg;
} ZoneState;

void zn_free(Zone *z);
void zn_reset(Zone *z);
void zn_save(Zone *z, ZoneState *m);
void zn_load(Zone *z, ZoneState *m);

void *zn_alloc(Zone *z, ptrdiff_t n);
void *zn_alloc_align(Zone *z, ptrdiff_t n, size_t align);
void *zn_realloc(Zone *z, void *ptr, ptrdiff_t oldsz, ptrdiff_t newsz);
void *zn_realloc_align(Zone *z, void *ptr, ptrdiff_t oldsz, ptrdiff_t newsz, size_t align);

char *zn_strdup(Zone *z, const char *s);
void *zn_zeroed(void *, size_t);

#define zn_new(z, t) (t*)zn_zeroed(zn_alloc_align(z, sizeof(t), alignof(t)), sizeof(t))

#ifdef ZONE_IMPL

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "zone.h"

#ifndef ZONE_BACKEND
	#ifdef __unix__
		#define ZONE_BACKEND ZONE_USE_MMAP
	#else
		#define ZONE_BACKEND ZONE_USE_MALLOC
	#endif
#endif

#if ZONE_BACKEND == ZONE_USE_MMAP

	#include <sys/mman.h>
	#include <unistd.h>

	static ZoneFrame *zn_zf_allocate(size_t capacity) {
		ZoneFrame *zf = mmap(NULL, sizeof *zf + capacity * sizeof(uintptr_t), PROT_READ | PROT_WRITE, MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
		if (!zf) {
			fprintf(stderr, "failed to allocate memory zone frame\n");
			abort();
		}
		return zf;
	}
	static void zn_zf_free(ZoneFrame *zf) {
		munmap(zf, (uintptr_t)zf->end - (uintptr_t)zf);
	}
	static size_t zn_zf_capacity(size_t cap) {
		if (cap < ZONE_MIN_CAPACITY) cap = ZONE_MIN_CAPACITY;
		size_t pg_size = sysconf(_SC_PAGE_SIZE) / sizeof(uintptr_t);
		size_t hdr_size = sizeof(ZoneFrame) / sizeof(uintptr_t);
		return cap + (-cap & (pg_size - 1)) - hdr_size;
	}

#elif ZONE_BACKEND == ZONE_USE_MALLOC

	static ZoneFrame *zn_zf_allocate(size_t capacity) {
		ZoneFrame *zf = malloc(sizeof *zf + capacity * sizeof(uintptr_t));
		if (!zf) {
			fprintf(stderr, "failed to allocate memory zone frame\n");
			abort();
		}
		return zf;
	}
	static void zn_zf_free(ZoneFrame *zf) {
		free(zf);
	}
	static size_t zn_zf_capacity(size_t cap) {
		if (cap < ZONE_MIN_CAPACITY) cap = ZONE_MIN_CAPACITY;
		return cap;
	}

#else

	#error "unknown or unsupported zone backend"

#endif

static ZoneFrame *zn_zf_new(size_t capacity) {
	capacity = zn_zf_capacity(capacity);
	ZoneFrame *zf = zn_zf_allocate(capacity);
	zf->prev = NULL;
	zf->next = NULL;
	zf->beg = zf->data;
	zf->end = (uint8_t*)zf + capacity;
	return zf;
}

#define PTR_ALIGN(ptr, align) (typeof(ptr))((uint8_t*)ptr + ((-(uintptr_t)ptr) & ((align)-1)))

void *zn_alloc_align(Zone *z, ptrdiff_t n, size_t align) {
	ZoneFrame *zf = z->cur;
	while (zf && PTR_ALIGN(zf->beg, align) + n >= zf->end) {
		zf = zf->next;
	}
	if (!zf) {
		zf = zn_zf_new(n + align - 1);
		zf->prev = z->tail;
		if (z->tail) z->tail->next = zf;
		z->tail = zf;
		z->cur = zf;
	}
	zf->beg = PTR_ALIGN(zf->beg, align);
	void *p = zf->beg;
	zf->beg += n;
	return p;
}

void *zn_realloc_align(Zone *z, void *ptr, ptrdiff_t oldsz, ptrdiff_t newsz, size_t align) {
	goto justalloc;
	if (!ptr || !oldsz) return zn_alloc_align(z, newsz, align);
	if (z->cur && ptr == z->cur->beg - oldsz && z->cur->beg - oldsz + newsz < z->cur->end) {
		z->cur->beg -= oldsz;
		z->cur->beg += newsz;
		return ptr;
	} else {
justalloc:
		void *p = zn_alloc_align(z, newsz, align);
		memcpy(p, ptr, oldsz);
		return p;
	}
}

static inline size_t zn_align_for(size_t n) {
	size_t a = 1;
	while (n > a) a <<= 1;
	return a;
}

void *zn_alloc(Zone *z, ptrdiff_t n) {
	return zn_alloc_align(z, n, zn_align_for(n));
}

void *zn_realloc(Zone *z, void *ptr, ptrdiff_t oldsz, ptrdiff_t newsz) {
	return zn_realloc_align(z, ptr, oldsz, newsz, zn_align_for(newsz));
}

void *zn_calloc(Zone *z, size_t n, size_t size) {
	void *p = zn_alloc_align(z, n * size, zn_align_for(size));
	memset(p, 0, n * size);
	return p;
}

void zn_free(Zone *z) {
	ZoneFrame *a = z->tail, *b;
	while (a) {
		b = a->prev;
		zn_zf_free(a);
		a = b;
	}
}

void zn_reset(Zone *z) {
	ZoneFrame *zf = z->tail;
	while (zf) {
		zf->beg = zf->data;
		if (!zf->prev) break;
		zf = zf->prev;
	}
	if (zf) z->cur = zf;
}

void zn_save(Zone *z, ZoneState *m) {
	m->zf = z->cur;
	m->beg = z->cur ? z->cur->beg : 0;
}

void zn_load(Zone *z, ZoneState *m) {
	if (m->zf) {
		while (z->cur != m->zf) {
			z->cur->beg = z->cur->data;
			z->cur = z->cur->prev;
		}
		m->zf->beg = m->beg;
	} else {
		zn_reset(z);
	}
}

/* utils */

char *zn_strdup(Zone *z, const char *s) {
	size_t n = strlen(s) + 1;
	char *d = zn_alloc_align(z, n, 1);
	memcpy(d, s, n);
	return d;
}

void *zn_zeroed(void *ptr, size_t n) {
	memset(ptr, 0, n);
	return ptr;
}

#endif
#endif