#include <stdio.h>
#include <string.h>
#include <assert.h>
#include "buddy.h"

static uint8_t MAGIC[MAGIC_COOKIE_BYTES];

struct free_list {
    uint8_t magic[MAGIC_COOKIE_BYTES];
    struct free_list *next;
    struct free_list **pprev;
};

static struct free_list *(AVAIL[ADDRESS_BITS]);
static size_t ROOT_LOGSIZE;
static void *BASE;

// log2 of size rounded up (ceil) or down (!ceil) to the nearest power of 2
static size_t size2log(size_t size, int ceil) {
    int bitcnt = 0;
    for (int i = 0; size; i++) {
        bitcnt += size & 1;
        size >>= 1;
        if (size) continue;
        if (ceil) return (bitcnt > 1) ? (i + 1) : i;
        return i;
    }
    return ceil ? 1 : 0;
}

void init_buddy(uint8_t *base, size_t size,
                uint8_t magic[MAGIC_COOKIE_BYTES]) {
    BASE = base;

    size_t logsize = size2log(size, 0);
    ROOT_LOGSIZE = logsize;

    AVAIL[logsize] = (struct free_list *)base;
    AVAIL[logsize]->next = 0;
    AVAIL[logsize]->pprev = &(AVAIL[logsize]);
    memcpy(MAGIC, magic, sizeof(MAGIC));
}

static struct free_list *buddy_of(struct free_list *block, size_t logsize) {
    size_t virt_block = (size_t)block - (size_t)BASE;
    size_t virt_buddy = virt_block ^ (1 << logsize);
    return (void*)((uint8_t*)BASE + virt_buddy);
}

static int isfree(struct free_list *block) {
    return !memcmp(block->magic, MAGIC, sizeof(MAGIC));
}

static void makefree(struct free_list *block) {
    memcpy(block->magic, MAGIC, sizeof(MAGIC));
}

static struct free_list *pop(struct free_list *block) {
    *(block->pprev) = block->next;
    if (block->next) block->next->pprev = block->pprev;
    return block;
}

static void push(struct free_list *block, size_t logsize) {
    block->next = AVAIL[logsize];
    if (block->next) block->next->pprev = &(block->next);
    AVAIL[logsize] = block;
    block->pprev = &(AVAIL[logsize]);
}

static void *_allocate(size_t logsize) {
    if (logsize > ROOT_LOGSIZE) return 0;
    if (AVAIL[logsize]) {
        struct free_list *block = pop(AVAIL[logsize]);
        memset(block, 0, sizeof(struct free_list));
        return block;
    }

    struct free_list *parent = _allocate(logsize + 1);
    struct free_list *buddy = buddy_of(parent, logsize);
    // split @parent in half and place the buddy on the avail list.
    memcpy(buddy->magic, MAGIC, sizeof(MAGIC));
    push(buddy, logsize);
    return parent;
}
void *allocate(size_t size) {
    size = (size < sizeof(struct free_list)) ? sizeof(struct free_list) : size;
    return _allocate(size2log(size, 1));
}

static void _liberate(struct free_list *block, size_t logsize) {
    push(block, logsize);
    if (logsize == ROOT_LOGSIZE) return;

    struct free_list *buddy = buddy_of(block, logsize);
    if (!isfree(buddy)) return;

    // coalesce up!
    struct free_list *smaller = (buddy < block) ? buddy : block;
    struct free_list *bigger = (buddy < block) ? block : buddy;
    pop(bigger);
    memset(bigger, 0, sizeof(*bigger));

    pop(smaller);

    _liberate(smaller, logsize + 1);
}

void liberate(void *base, size_t size) {
    struct free_list *block = base;
    memset(block, 0, sizeof(*block));
    makefree(block);

    _liberate(block, size2log(size, 1));
}

void debug_buddy(void) {
    for (size_t i = 0; i <= ROOT_LOGSIZE; i++) {
        printf("Free blocks at size 2^%lu = %lu:\n", i, 1ul << i);
        for (struct free_list *block = AVAIL[i]; block; block = block->next) {
            printf("\t%p\n", block);
        }
    }
}