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

static uint8_t MAGIC[MAGIC_COOKIE_BYTES];

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

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));
}

// interpret @block as a block at depth @logsize. is it free?
static int isfree_at_logsize(struct free_list *block, size_t logsize) {
    return isfree(block) && block->logsize == logsize;
}

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]);
    block->logsize = 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_at_logsize(buddy, logsize)) 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);
        }
    }
}

///////// "advanced features"
static struct free_list *rhs_child_of(struct free_list *block, size_t logsize) {
    size_t virt_block = (size_t)block - (size_t)BASE;
    size_t virt_child = virt_block | (1 << (logsize - 1));
    return (void*)((uint8_t*)BASE + virt_child);
}

// NOTE: this method is perhaps more complicated than it needs to be because we
// take great pains to avoid writing to the region that is being reserved
// (e.g., in case it is device MMIO).
int reserve(void *start, size_t size) {
    assert(size);
    void *end = (void*)((uint8_t*)start + size);

    // (1) iterate down to find the first free node to the left of @start
    struct free_list *parent = (void*)BASE;
    size_t parent_logsize = ROOT_LOGSIZE;
    while (!isfree_at_logsize(parent, parent_logsize)) {
        // pick the proper child and recurse
        struct free_list *rhs_child = rhs_child_of(parent, parent_logsize);
        if ((void*)rhs_child <= start) parent = rhs_child;
        if (!parent_logsize--) return 0;
    }
    // parent is free!
    void *parent_end = (void*)((uint8_t*)parent + (1 << parent_logsize));
    // if the size has no chance of fitting, return 0.
    if (parent_end < end) return 0;
    // otherwise, split this parent
    pop(parent);
    while (1) {
        // check whether the child could fit it.
        struct free_list *rhs_child = rhs_child_of(parent, parent_logsize);
        struct free_list *next_child = ((void*)rhs_child <= start) ? parent : rhs_child;
        void *child_end = (void*)((uint8_t*)next_child + (1 << (parent_logsize - 1)));
        if (child_end < end) {
            // the child *cannot* fit it, so allocate this parent.
            return 1;
        } else {
            // the child *can* fit it, so split this parent into two children.
            struct free_list *lhs_child = parent;
            if ((void*)rhs_child <= start) {
                // the lhs child will be free
                makefree(lhs_child);
                push(lhs_child, parent_logsize - 1);
                parent = rhs_child;
            } else {
                // the rhs child will be free
                makefree(rhs_child);
                push(rhs_child, parent_logsize - 1);
                parent = lhs_child;
            }
        }
        parent_logsize--;
    }
    return 0;
}

static void *_naive_reallocate(void *old, size_t old_size, size_t new_size) {
    void *new = allocate(new_size);
    if (!new) return 0;
    assert(old_size < new_size);
    memcpy(new, old, old_size);
    liberate(old, old_size);
    return new;
}

void *reallocate(void *old, size_t old_size, size_t new_size) {
    size_t old_logsize = size2log(old_size, 1);
    size_t new_logsize = size2log(new_size, 1);

    if (new_logsize == old_logsize) return old;

    if (new_logsize < old_logsize) {
        // repeatedly split, keeping lhs.
        struct free_list *block = old;
        while (new_logsize < old_logsize) {
            old_logsize--;
            struct free_list *right_half = buddy_of(block, old_logsize);
            makefree(right_half);
            push(right_half, old_logsize);
        }
        return old;
    }

    // otherwise, we must iterate up the tree and ensure that at each level:
    // (1) we are the left-buddy, and
    // (2) our buddy is free.
    // up until we reach an ancestor of the proper size.

    // First, just verify this claim.
    size_t pos_logsize = old_logsize;
    if (new_logsize > ROOT_LOGSIZE) return 0;
    struct free_list *pos = old;
    while (pos_logsize != new_logsize) {
        struct free_list *buddy = buddy_of(pos, pos_logsize);
        if (pos > buddy) {
            // oh no, we're the right buddy!
            return _naive_reallocate(old, old_size, new_size);
        } else if (!isfree_at_logsize(buddy, pos_logsize)) {
            // oh no, our buddy at this level isn't free!
            return _naive_reallocate(old, old_size, new_size);
        }
        pos_logsize++;
    }

    // Then, revisit the path and coalesce on the way up.
    pos_logsize = old_logsize;
    pos = old;
    while (pos_logsize != new_logsize) {
        struct free_list *buddy = buddy_of(pos, pos_logsize);
        pop(buddy);
        memset(buddy, 0, sizeof(struct free_list));
        pos_logsize++;
    }
    return old;
}