/* Simple heap allocator.
 *
 * I saw this article: https://samwho.dev/memory-allocation/
 * and was inspired to re-write the heap allocator I made a
 * while ago to try to make it faster and less wasteful.
 *
 * How it works:
 * The allocator is given a buffer where all of the allocations
 * will go. Each allocation is given a "block" from the buffer
 * that is the size of the allocation plus a "header size". The
 * header of each block is four bytes: The most significant
 * 31 bits represent the size of the allocation and the least
 * significant bit (or the "taken" bit) is 1 if the block
 * currently represents an allocation and 0 if it's free and
 * available for use.
 *
 * On initialisation, the allocator is given a single free
 * block that is the size of the entire buffer.
 *
 * When heap_alloc is called:
 *   First, all of the current blocks are iterated. If a free
 *   block is found  that's the same size as the requested
 *   allocation, then the block's "taken" bit is set to 1 and
 *   the block's allocation is returned.
 *   Otherwise, a free block that is larger than the requested
 *   allocation may be found. In this case, the block is split
 *   in two such that two blocks remain: A block of the same
 *   size as the allocation and a block containing the remaining
 *   half of the block. The former is returned. The free block
 *   from the split is placed before the returned block (in
 *   other words, the buffer is filled from back to front),
 *   that way when the allocator is called on again it might
 *   well find a suitable block within a few iterations.
 *   Otherwise, no free blocks were found to be big enough
 *   and thus the program is out of memory. Here, it just calls
 *   abort, but this case may want handling more cleanly.
 *
 * When heap_free is called:
 *   The header to the allocation is first calculated - it's
 *   simply the four bytes that come directly before the
 *   pointer that is given. The header's "taken" bit is set
 *   to zero to indicate a freed block.
 *   Then, the allocator is de-fragmented. What this boils
 *   down to is iterating all of the blocks and recording the
 *   first free block that is seen. Then, all free blocks after
 *   that are iterated until a "taken" block is found. Every
 *   block between the first free block and the first "taken"
 *   block is coagulated into a single free block before
 *   continuing from the first step until every block has been
 *   visited.
 *
 * There are also functions for aligning allocations, as
 * that is also an integral part of any allocator. By default,
 * allocations are given an alignment of 4 bytes, which is
 * typical of general-purpose allocators on PCs.
 */

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

typedef struct {
	char* buffer;
	int size, blocks;
} Heap;

void init_heap(Heap* heap, char* buf, int size) {
	heap->buffer = buf;
	heap->size = size;
	heap->blocks = 1;
	((int*)buf)[0] = size << 1;
}

static uintptr_t align_address(
	uintptr_t address,
	size_t align
) {
	size_t m;
	m = align - 1;
	return (address + m) & ~m;
}

static void* heap_alloc_impl(Heap* h, int size) {
	int* header, * nh, s, f, i, as, p;
	i = 0;
	as = size + sizeof *header;
	p = 0;
	for (i = 0; i < h->blocks; i++) {
		header = (int*)(h->buffer + p);
		s = header[0];
		f = !(s & 1);
		s >>= 1;
		if (f) {
			/* There's no point creating blocks smaller than
			 * the header. */
			if (as < s - sizeof *header - 1) {
				s -= as;
				header[0] = (s << 1);
				nh = (int*)(h->buffer + p + s);
				nh[0] = (as << 1) | 1;
				h->blocks++;
				return nh + 1;
			} else if (s == as) {
				header[0] |= 1;
				return header + 1;
			}
		}
		p += s;
	}
	abort();
	return 0;
}

static void heap_free_impl(Heap* h, void* ptr) {
	int* header, *h2, i, j, s, f, d, n, ns, p;
	header = &((int*)ptr)[-1];
	*header &= INT32_MAX << 1;
	p = n = 0;
	for (i = 0; i < h->blocks; i++) {
		header = (int*)(h->buffer + p);
		s = header[0];
		f = !(s & 1);
		s >>= 1;
		if (f) {
			ns = s;
			p += s;
			d = 0;
			for (j = i + 1; j < h->blocks; j++) {
				h2 = (int*)(h->buffer + p);
				s = h2[0];
				f = !(s & 1);
				if (!f) { break; }
				s >>= 1;
				p += s;
				ns += s;
				d++;
			}
			i = j - 1;
			n += d;
			if (d)
				header[0] = ns << 1;
		} else
			p += s;
	}
	h->blocks -= n;
}

void heap_print_blocks(const Heap* h) {
	const int* header;
	int s, f, p, i;
	p = 0;
	printf("%8s %8s %s\n", "offset", "size", "free");
	for (i = 0; i < h->blocks; i++) {
		header = (const int*)(h->buffer + p);
		s = header[0];
		f = !(s & 1);
		s >>= 1;
		printf("%8d %8d %d\n", p, s, f);
		p += s;
	}
}

void* heap_alloc_aligned(Heap* h, int size, unsigned align) {
	unsigned char* p, * a;
	ptrdiff_t shift;
	size += (int)align;
	p = heap_alloc_impl(h, size);
	a = (unsigned char*)align_address((uintptr_t)p, align);
	a += align * (unsigned)(p == a);
	shift = a - p;
	a[-1] = shift & 0xff;
	return a;
}

void heap_free_aligned(Heap* h, void* ptr) {
	unsigned char* a;
	ptrdiff_t shift;
	a = ptr;
	shift = a[-1];
	shift += 256 * shift == 0;
	a -= shift;
	heap_free_impl(h, a);
}

void* heap_alloc(Heap* h, int size) {
	return heap_alloc_aligned(h, size, 4);
}

void heap_free(Heap* h, void* ptr) {
	heap_free_aligned(h, ptr);
}

int main() {
	char buffer[4096];
	Heap heap;
	int i, s;
	void* allocs[64];
	init_heap(&heap, buffer, sizeof buffer);
	srand(500);
	/* Test: allocate a bunch of random blocks,
	 * then free them randomly, then allocate some
	 * more and print the state. */
	for (i = 0; i < 64; i++) {
		s = rand() % 32 + 1;
		allocs[i] = heap_alloc(&heap, s);
	}
	for (i = 0; i < 64; i++) {
		if (rand() % 5 == 0) {
			heap_free(&heap, allocs[i]);
		}
	}
	for (i = 0; i < 32; i++) {
		s = rand() % 32 + 1;
		heap_alloc(&heap, s);
	}
	heap_print_blocks(&heap);
}