/*
 * Toy reimplementation of malloc/free, unoptimized. Did this when got bored waiting for plane at airport.
 */

#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <assert.h>
#include <stdlib.h>

/*
 * struct used to store allocated segments
 */
struct segment {
    void *start;
    size_t size;
    struct segment *next;
} /*__attribute__((packed))*/;

struct segment *__segment_list;     /* start of allocated segments circular list */
void *initial_brk;                  /* initial brk stored at beginning of the program */

struct segment *lookup_segment(void *addr)
{
    struct segment *s = __segment_list;
    do {
        if (s->start == addr)
            return s;
        s = s->next;
    } while (s != __segment_list);

    return NULL;
}

/* this function is horrible, wouldn't need it if I was using double linked list */
struct segment *lookup_segment_before(struct segment *s)
{
    struct segment *s2 = __segment_list;
    do {
        if (s2->next == s)
            return s2;
        s2 = s2->next;
    } while (s2 != __segment_list);

    return NULL;
}

void insert_segment(struct segment *s)
{
    struct segment *s2 = __segment_list;

    if (!__segment_list) {
        s->next = s;
        __segment_list = s;
    } else {
        do {
            if (s->start > s2->start && s->start < s2->next->start)
                break;
            s2 = s2->next;
        } while (s2 != __segment_list);

        s->next = s2->next;
        s2->next = s;
    }
}

void *__get_chunk(size_t size)
{
    struct segment *s = __segment_list;
    void *addr = NULL;

    if (!s)
        return NULL;

    if ((unsigned)((s->start - sizeof(*s)) - initial_brk) >= size)
        return initial_brk;

    do {
        if (s != s->next && (unsigned)((s->next->start - sizeof(*s)) - (s->start + s->size)) >= size) {
            addr = s->start + s->size;
            break;
        }

        s = s->next;
    } while (s != __segment_list);

    return addr;
}

/*
 * do the effective work - increase brk
 */
void *__my_malloc(size_t size)
{
    void *ret;
    ret = sbrk(size);
    if (ret == (void *)-1) {
        perror("sbrk");
        return NULL;
    }

    return ret;
}

/*
 * create a new struct segment, add it to the circular list
 * and call __my_malloc()
 */
void *my_malloc(size_t size)
{
    struct segment *s;
    void *start, *ret, *chunk;

    if ((chunk = __get_chunk(sizeof(*s) + size)) != NULL) {
        printf("found chunk at %p\n", chunk);
        s = chunk;
        start = ret = chunk+sizeof(*s);
    } else {
        s = __my_malloc(sizeof(*s));
        start = ret = __my_malloc(size);
    }

    assert(ret+size <= sbrk(0));
    if (!ret)
        return NULL;

    s->start = start;
    s->size = size;

    insert_segment(s);

    return ret;
}

/*
 * calculate the highest brk address and adjust it if needed
 */
void __update_brk()
{
    void *max_addr = initial_brk;
    struct segment *s = __segment_list;
    int ret;

    if (!s)
        goto setbrk;

    do {
        if (s->start + s->size > max_addr)
            max_addr = s->start + s->size;
        s = s->next;
    } while (s != __segment_list);

    if (!max_addr)
        abort();

setbrk:
    ret = brk(max_addr);
    if (ret) {
        perror("brk");
        abort();
    }
}

void my_free(void *addr)
{
    struct segment *s2, *s = lookup_segment(addr);
    if (!s)
        abort();

    s2 = lookup_segment_before(s);
    if (!s2)
        abort();

    if (s == __segment_list) {
        __segment_list = s->next;
        if (__segment_list->next == __segment_list)
            __segment_list = NULL;
    }

    s2->next = s->next;

    __update_brk();
}

void print_brk(char *msg)
{
    printf("brk: %p (%s)\n", sbrk(0), msg);
}

int main()
{
    initial_brk = sbrk(0);

    print_brk("initial brk");
    char *test = my_malloc(16*sizeof(char));
    print_brk("first alloc");
    char *test2 = my_malloc(16*sizeof(char));
    print_brk("second alloc");
    my_free(test);
    print_brk("free of first");
    test = my_malloc(16*sizeof(char));
    print_brk("third alloc");
    my_free(test);
    print_brk("free of third");
    my_free(test2);
    print_brk("free of second");
    return 0;
}