/* File released under CC0 (public domain), read more
   https://creativecommons.org/publicdomain/zero/1.0/deed.en
   */

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

#define HEAP_MAGIC 0x1DEFEC8D
#define HEAP_SIZE 0x100000
#define HEAP_FLAG_FREE (1 << 0)

typedef struct heap_header{
    uint32_t magic;
    uint32_t flags;
    uint64_t size;
    struct heap_header* prev;
    struct heap_header* next;
} heap_header_t;

static char heap[HEAP_SIZE] ;
static heap_header_t* heap_start;

static inline char is_block_free(heap_header_t* block)
{
    if(block == NULL) return 0;
    return (block->flags & HEAP_FLAG_FREE);
}

static void dump_header(heap_header_t* header);
static void* initialize_heap(uint64_t size);

static inline heap_header_t* find_free_block(uint64_t size);
static inline void split_block(heap_header_t* block, uint64_t size);
static inline void* allocate_block(uint64_t size);

static inline void deallocate_block(heap_header_t* address);

void* malloc(size_t size);
void free(void* ptr);

int main(void)
{
    heap_start = initialize_heap(HEAP_SIZE);

    dump_header(allocate_block(0x40));
    dump_header(allocate_block(0x40));
    dump_header(allocate_block(0x40));

    printf("%p, %p, %p\n", malloc(0x20), malloc(0x500), malloc(0x800));
    void* ptr1 = malloc(0x100);
    void* ptr2 = malloc(0x100);
    void* ptr3 = malloc(0x100);
    printf("%p, %p, %p\n", malloc(0x20), malloc(0x500), malloc(0x800));

    free(ptr2);
    dump_header((heap_header_t*)((uint64_t)ptr1 - sizeof(heap_header_t)));
    dump_header((heap_header_t*)((uint64_t)ptr2 - sizeof(heap_header_t)));
    dump_header((heap_header_t*)((uint64_t)ptr3 - sizeof(heap_header_t)));
    free(ptr1);
    dump_header((heap_header_t*)((uint64_t)ptr1 - sizeof(heap_header_t)));
    dump_header((heap_header_t*)((uint64_t)ptr2 - sizeof(heap_header_t)));
    dump_header((heap_header_t*)((uint64_t)ptr3 - sizeof(heap_header_t)));

    return 0;
}

void dump_header(heap_header_t* header)
{
    if(header != NULL)
        printf("%p, 0x%X, 0x%lX, %p, %p\n",
                header,
                header->flags,
                header->size,
                header->prev,
                header->next
              );
    else
        printf("(nil)\n");
}

void* initialize_heap(uint64_t size)
{
    heap_header_t* header = (heap_header_t*)&heap[0];

    header->magic = HEAP_MAGIC;
    header->flags = HEAP_FLAG_FREE;
    header->size = size - sizeof(*header);
    header->prev = NULL;
    header->next = NULL;

    return (void*)header;
}

heap_header_t* find_free_block(uint64_t size)
{
    heap_header_t* current = (heap_header_t*)&heap[0];

    while(current != NULL)
    {
        if(current->size >= size && is_block_free(current))
            return current;

        current = current->next;
    }

    return current;
}

void split_block(heap_header_t* block, uint64_t size)
{
    heap_header_t* new_block =
        (heap_header_t*)((uint64_t)block + sizeof(*block) + size);

    new_block->magic = HEAP_MAGIC;
    new_block->flags = HEAP_FLAG_FREE;
    new_block->size = block->size - size - sizeof(*new_block);
    new_block->prev = block;
    new_block->next = block->next;

    block->size = size;
    block->next = new_block;
}

void* allocate_block(uint64_t size)
{
    heap_header_t* block = find_free_block(size);

    if(block == NULL) return NULL;

    if(block->size - size > sizeof(*block))
        split_block(block, size);

    block->flags &= (uint32_t)~HEAP_FLAG_FREE;

    return block;
}

void deallocate_block(heap_header_t* block)
{
    block->flags |= (uint32_t)HEAP_FLAG_FREE;
    while(is_block_free(block->next))
        block->next = block->next->next;
    while(is_block_free(block->prev))
        block->prev = block->prev->prev;

    block->next->prev = block->prev;
}

void* malloc(size_t size)
{
    heap_header_t* block = allocate_block((uint64_t)size);

    if(block == NULL) return NULL;

    return (void*)((uint64_t)block + sizeof(*block));
}

void free(void* ptr)
{
    if(ptr == NULL) return;

    heap_header_t* block = (heap_header_t*)((uint64_t)ptr - sizeof(*block));
    deallocate_block(block);
}