#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* 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->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;
}

/// 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
		 	Fragment here if possible.
		*/
	}
	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.
/// @todo Reverse defragmenting.
void free(void* ptr)
{
	if (!ptr)
		return;

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

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

	// Reverse defragment here.

	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", *b);

	free(b);
	free(a);

	return 0;
}