#define _BSD_SOURCE
#include <stdint.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <sys/mman.h>

#define u16set(hi, lo) ((hi << 8) | (lo & 0xff))
#define align_up(n, align) (((n) + (align - 1)) & (-align))
#define header_set_magic(block, mag, poi) ((block)->magic = u16set(poi, mag))
#define Assert(pred) if (!(pred)) {__builtin_trap();}

#define BUDDY_BLOCK_MAGIC 0xb7
#define BUDDY_BLOCK_FREED 0xaa
#define BUDDY_MIN_ALIGNMENT sizeof(BuddyBlockHeader)

#define false ((bool)0)
#define true ((bool)1)
#define nil ((void *)0)

typedef uint8_t u8;
typedef _Bool bool;
typedef ptrdiff_t isize;

/* Header of a buddy block. This gets baked into memory. */
typedef struct {
	isize size;
	bool used;
	uint16_t magic;
} BuddyBlockHeader;
/* Being 16 bytes the minimum alignment, at least on a 64-bit machine... */
_Static_assert(sizeof(BuddyBlockHeader) == 16, "BuddyBlockHeader size is not 16 bytes");

typedef struct {
	BuddyBlockHeader *head, *tail;	
	isize backing_len;
} BuddyAllocator;

static void
abort_msg(const char *s)
{
	fputs(s, stderr);
	fflush(stderr);
	abort();
}

static BuddyBlockHeader *
next_buddy_block(BuddyBlockHeader *cur)
{
	/* Calculate the address for the next buddy block header start.
	 * This has to be padded to a boundary of at least
	 * `sizeof(BuddyBlockHeader)` bytes. */
	return (BuddyBlockHeader *)align_up(
		(uintptr_t)((u8 *)cur + cur->size), BUDDY_MIN_ALIGNMENT);
}

/* Split a block until it has a size equal or greater than `size`.
 * This function also "builds" the buddy headers. */
static BuddyBlockHeader *
split_block(BuddyBlockHeader *block, isize size)
{
	BuddyBlockHeader *cur = block;
	while (cur->size > size) {
		/* XXX: Splitting by half seems quite wasteful... (see `buddy_alloc_aligned`) */
		const isize half = cur->size / 2;
		cur->size = half;
		cur = next_buddy_block(cur);
		cur->size = half;
		cur->used = false;
		cur->magic = u16set(0, BUDDY_BLOCK_MAGIC);
	}
	Assert(cur->size >= size);
	return cur;
}

static BuddyBlockHeader *
find_block(BuddyAllocator *ba, isize size)
{
	BuddyBlockHeader *cur = ba->head;
	BuddyBlockHeader *next = next_buddy_block(cur);
	/* Case for an initial state, where there aren't any split buddies */
	if (!cur->used && next == ba->tail)
		return split_block(cur, size);

	/* Traverse the block list for a suitable block */
	BuddyBlockHeader *block = nil;
	while (cur < ba->tail) {
		Assert((cur->magic & 0xff) == BUDDY_BLOCK_MAGIC);
		Assert(cur->size <= ba->backing_len);
		/* Coalesce two adjacent buddies if possible */
		if (next < ba->tail
			&& !next->used
			&& next->size == cur->size
			&& next->size + cur->size >= size) {
			Assert((next->magic & 0xff) == BUDDY_BLOCK_MAGIC);
			fprintf(stderr, "candidate for coalescing (%p U %p)\n", (void *)cur, (void *)next);
			cur->size += next->size;
			/* wipe header of the buddy just in case */
			memset(next, 0, sizeof(*next));
		}
		if (!cur->used && cur->size >= size) {
			fprintf(stderr, "candidate at %p (size: %td)\n", (void *)cur, cur->size);
			/* Either if no candidate has been picked yet or the current candidate
			 * is smaller (but `size` can still fit) than the current buddy. */
			if (block == nil || block->size > cur->size)
				block = cur;
		}
		cur = next_buddy_block(cur);
		next = cur < ba->tail ? next_buddy_block(cur) : nil;
	}
	if (block != nil)
		return split_block(block, size);
	return nil;
}

static void
coalesce_buddies(BuddyAllocator *ba)
{
	BuddyBlockHeader *cur = ba->head;
	BuddyBlockHeader *next = next_buddy_block(cur);

	while (cur < ba->tail) {
		Assert((cur->magic & 0xff) == BUDDY_BLOCK_MAGIC);
		Assert(cur->size <= ba->backing_len);
		if (next < ba->tail
			&& !next->used
			&& next->size == cur->size) {
			Assert((next->magic & 0xff) == BUDDY_BLOCK_MAGIC);
			cur->size += next->size;
			memset(next, 0, sizeof(*next));
		}
		cur = next_buddy_block(cur);
		next = cur < ba->tail ? next_buddy_block(cur) : nil;
	}
}

static void
buddy_alloc_init(BuddyAllocator *balloc, u8 *buf, isize buflen)
{
	balloc->head = (BuddyBlockHeader *)buf;
	balloc->tail = (BuddyBlockHeader *)(buf + buflen);
	balloc->backing_len = buflen;
	*balloc->head = (BuddyBlockHeader) {
		.size = buflen,
		.used = false,
		.magic = u16set(0, BUDDY_BLOCK_MAGIC),
	};
}

void *
buddy_alloc_aligned(BuddyAllocator *ba, isize size, isize alignment)
{
	if (size == 0)
		return nil;

	Assert(size > 0);

	if (alignment < BUDDY_MIN_ALIGNMENT)
		alignment = BUDDY_MIN_ALIGNMENT;
	isize aligned_size = align_up(size + sizeof(BuddyBlockHeader), alignment);
	isize ss = alignment; /* squared up size (because we split the buddies in half) */
	while (ss <= aligned_size)
		ss *= 2;
	Assert(ss <= ba->backing_len);

	BuddyBlockHeader *block;
	if ((block = find_block(ba, ss)) != nil)
		goto happy;
	/* First attempt failed, maybe due to fragmentation. Do a run of coalescing to
	 * solve fragmentation if any and try allocating again. */
	coalesce_buddies(ba);
	if ((block = find_block(ba, ss)) != nil)
		goto happy;
	/* Too fragmented or ran out of memory. */
	return nil;
happy:
	block->used = true;
	block->magic &= 0xff; /* clear free'd flag */
	return (u8 *)block + sizeof(BuddyBlockHeader);
}

void *
buddy_alloc(BuddyAllocator *ba, isize size)
{
	return buddy_alloc_aligned(ba, size, BUDDY_MIN_ALIGNMENT);
}

void
buddy_free(BuddyAllocator *ba, void *ptr)
{
	if (ptr == nil)
		return;

	BuddyBlockHeader *header = (BuddyBlockHeader *)((u8 *)ptr - sizeof(*header));
	Assert((header->magic & 0xff) == BUDDY_BLOCK_MAGIC);
	/* double free detection */
	Assert((header->magic >> 8) != BUDDY_BLOCK_FREED);
	header->used = false;
	header->magic = u16set(BUDDY_BLOCK_FREED, BUDDY_BLOCK_MAGIC);
}

static void
iter_blocks(BuddyBlockHeader *head, BuddyBlockHeader *tail)
{
	int cc = fprintf(stderr, "------ list start (head: %p; tail: %p; size: %td) ------\n",
					(void *)head, (void *)tail, tail - head);

	BuddyBlockHeader *cur = head;
	while (cur < tail) {
		Assert((cur->magic & 0xff) == BUDDY_BLOCK_MAGIC);
		fprintf(stderr, "block@%p(size: %8td; used: %3s; magic: %04x)\n",
				(void *)cur, cur->size, cur->used? "yes":"no", cur->magic);
		cur = next_buddy_block(cur);
	}

	while (cc--)
		fputc('-', stderr);
	fputc('\n', stderr);
}

#include <time.h>

static void
dirt_block(u8 *buf, isize size)
{
	u8 *bufend = buf + size;
	while (buf < bufend)
		*buf++ = rand() & 0xff;
}

int
main(void)
{
	const isize backing_len = align_up(1 << 20, 16);
	u8 *backing = mmap(
		nil, backing_len,
		PROT_READ | PROT_WRITE,
		MAP_PRIVATE | MAP_ANONYMOUS,
		-1, 0
	);
	if (backing == MAP_FAILED) {
		perror("mmap");
		return 1;
	}
	*(volatile u8 *)backing = 0xaa; /* page fault it */

	BuddyAllocator ba = {0};
	buddy_alloc_init(&ba, backing, backing_len);

	srand(time(nil));

	iter_blocks(ba.head, ba.tail);
	u8 *buf1 = buddy_alloc(&ba, 255);
	dirt_block(buf1, 255);
	u8 *buf2 = buddy_alloc(&ba, 1024);
	dirt_block(buf2, 1024);
	buddy_free(&ba, buf1);
	iter_blocks(ba.head, ba.tail);
	u8 *buf3 = buddy_alloc(&ba, 512);
	dirt_block(buf3, 512);
	iter_blocks(ba.head, ba.tail);
	buddy_free(&ba, buf2);
	buddy_free(&ba, buf3);
	munmap(backing, backing_len);

	return 0;
}