#include <stddef.h>
#include <stdio.h>
#include <unistd.h>

#define ALIGN4(s) (((((s)-1)>>2)<<2)+4)
#define BLOCK_SIZE sizeof(struct block)

/// The memory block's header.
struct block
{
	size_t size;
	struct block* prev;
	struct block* next;
	int free;
};

struct block* first = NULL;
struct block* last = NULL;

/// Extends heap memory upwards, towards zero.
/// @param [in] s The size of the memory needed aligned by 4 bytes.
/// @returns The new memory block.
struct block* extend_heap(size_t s)
{
	struct block* b = sbrk(0);

	if (sbrk(BLOCK_SIZE + s) == (void*)-1)
		return NULL;

	b->size = s;
	b->prev = last;
	b->next = NULL;
	b->free = 0;

	if (last)
		last->next = b;

	last = b;

	return b;
}

/// Finds the first block that will fit the given size.
/// @param [in] s The 4 byte aligned size to look for.
/// @returns The matching available memory block.
struct block* find_first(size_t s)
{
	struct block* current = first;
	while (current && (!current->free || current->size < s))
		current = current->next;

	return current;
}

struct block* fragment_block(struct block* in, size_t s)
{
	size_t totalSize = BLOCK_SIZE + s;
	struct block* b = (struct block*)((char*)(in + 1) + (in->size - totalSize));
	b->size = s;
	b->prev = in;
	b->next = in->next;
	b->free = 0;

	in->size -= totalSize;
	in->next = b;

	return b;
}

/// Will find or allocate a memory block.
/// @param [in] size The size of the memory block to request.
/// @returns The requested memory on the heap.
/// @todo Fragmenting functionality.
void* malloc(size_t size)
{
	struct block* b;

	size_t alignedSize = ALIGN4(size);

	if (first)
	{
		b = find_first(alignedSize);
		if (!b)
			b = extend_heap(alignedSize);
		else if (b->size > BLOCK_SIZE + alignedSize)
			b = fragment_block(b, alignedSize);
	}
	else
	{
		b = extend_heap(alignedSize);
		if (!b)
			return NULL;

		first = b;
	}

	return b + 1;
}

/// Will flag the provided memory as free and will defragment other blocks adjacent to it.
/// @param [in] ptr The memory to flag as free.
/// @note If all data after the provided memory is free, it will reduce the heap size.
void free(void* ptr)
{
	if (!ptr)
		return;

	struct block* b = (struct block*)ptr - 1;
	if (b->free)
		return;

	b->free = 1;

	while (b->prev && b->prev->free)
		b = b->prev;

	while (b->next && b->next->free)
	{
		b->size += BLOCK_SIZE + b->next->size;
		b->next = b->next->next;
	}

	if (!b->next)
		brk(b);
}

int main()
{
	int* a = (int*)malloc(sizeof(int));
	int* b = (int*)malloc(sizeof(int));

	*a = 5;
	*b = 12;

	printf("Test 1: %i\n", *a);
	printf("Test 2: %i\n", *b);

	free(a);
	free(b);

	int* c = (int*)malloc(sizeof(int));

	return 0;
}