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

#include <assert.h>
#include <string.h>

#include <sys/syscall.h>
#include <unistd.h>

// TODO: do more error checking

#define BLOCK_SZ sizeof(struct_block)
#define BLOCK_ARR_CAP  1024

static inline void *_brk(void *addr) {
    return (void *)syscall(SYS_brk, addr);
}

typedef uint64_t u64;

struct block {
    void *addr;         // pointer to the start address of the block
    u64 size;           // the size of the block
    struct block *next; // pointer to the previous block
};

static void *heap_start;

static struct block *block_root;
static struct block block_arr[BLOCK_ARR_CAP];

void *mylloc(u64 size)
{
    if(heap_start == NULL) {
        heap_start = _brk(NULL);
    }

    void *addr = heap_start;

    struct block *next = block_root;
    struct block *prev = NULL;

    // find the first available spot to allocate
    while(next != NULL) {
        if(size <= (u64)(next->addr - addr)) {
            break;
        }

        addr += next->size;
        prev = next;
        next = next->next;
    }

    // no such spot, so extend the system break;
    if(next == NULL) {
        _brk(addr + size);
    }

    // find an empty block and  fill it with information
    for(int i = 0; i < BLOCK_ARR_CAP; i++) {
        if(block_arr[i].addr != NULL) {
            continue;
        }
        
        block_arr[i].addr = addr;
        block_arr[i].size = size;
        block_arr[i].next = next;

        if(prev == NULL) {
            block_root = &block_arr[i];
        } else {
            prev->next = &block_arr[i];
        }
        
        break;
    }

    return addr;
}

void myfree(void *addr)
{
    struct block *prev = NULL;
    struct block *block = block_root;
    
    while(block != NULL)
    {
        if(block->addr != addr) {
            prev = block;
            block = block->next;
            continue;
        }

        block->addr = NULL;

        if(prev == NULL) {
            block_root = block->next;
            if(block->next == NULL) {
                _brk(heap_start);
            }
        } else {
            prev->next = block->next;
            if(block->next == NULL) {
                _brk(prev->addr + prev->size);
            }
        }

        break;
    }
}

int main(void)
{
    char *str1 = mylloc(10);
    memset(str1, '2', 10);
    str1[9] = '\0';

    char *str2 = mylloc(20);
    memset(str2, '1', 20);
    str2[19] = '\0';

    myfree(str1);

    char *str3 = mylloc(11);
    memset(str3, '3', 11);
    str3[10] = '\0';

    printf("%s\n%s\n", str2, str3);

    myfree(str3);
    myfree(str2);
}