#include "buddymap.h"
#include "jove.h"
#include "print.h"

#define BLOCK_AT_LAYER(blocks, l, i) blocks[l][i / BUDDY_BLOCK_BITS]
#define BLOCK_BITMASK(i) (1ULL << (i % BUDDY_BLOCK_BITS))
#define LOG2(v) (31 - __builtin_clz(v))

bool
buddy_bit_test(struct BuddyMap *map, size_t l, size_t i)
{
    return (BLOCK_AT_LAYER(map->blocks, l, i) & BLOCK_BITMASK(i)) > 0;
}

void
buddy_bit_mark(struct BuddyMap *map, size_t l, size_t i)
{
    size_t wi = i << l;
    size_t bits = 1ULL << l;
    for(size_t wl = 0; wl < map->orders; wl++) {
        if(bits == 0) {
            if(buddy_bit_test(map, wl - 1, (wi << 1) + 1)) bits = 1;
        }
        for(size_t bit = 0; bit < bits; bit++) {
            size_t rbit = bit + wi;
            if(l == 0 && !buddy_bit_test(map, 0, rbit)) map->free--;
            BLOCK_AT_LAYER(map->blocks, wl, rbit) |= BLOCK_BITMASK(rbit);
        }
        bits >>= 1;
        wi >>= 1;
    }
}

void
buddy_bit_free(struct BuddyMap *map, size_t l, size_t i)
{
    size_t wi = i << l;
    size_t bits = 1ULL << l;
    for(size_t wl = 0; wl < map->orders; wl++) {
        if(bits == 0) bits = 1;
        for(size_t bit = 0; bit < bits; bit++) {
            size_t rbit = bit + wi;
            if(l == 0 && buddy_bit_test(map, 0, rbit)) map->free++;
            BLOCK_AT_LAYER(map->blocks, wl, rbit) &= ~BLOCK_BITMASK(rbit);
        }
        bits >>= 1;
        wi >>= 1;
    }
}

static void
s_buddy_op_range(
        struct BuddyMap *map,
        size_t l,
        size_t start,
        size_t end,
        void (*op)(struct BuddyMap*, size_t, size_t))
{
    size_t layerw = 1ULL << l;
    size_t start_off = start % layerw;
    size_t end_off = end % layerw;

    size_t start_real = start + (start_off > 0 ? (layerw - start_off) : 0);
    size_t end_real = end - end_off;

    if(start_real != start) s_buddy_op_range(map, l - 1, start, start_real, op);
    if(end_real != end) s_buddy_op_range(map, l - 1, end_real, end, op);

    size_t start_bit = start_real >> l;
    size_t end_bit = end_real >> l;
    if(start_bit == end_bit || end_bit < start_bit) return;
    for(size_t bit = start_bit; bit < end_bit; bit++)
        op(map, l, bit);
}

static size_t
s_buddy_layer_bestfit(struct BuddyMap *map, size_t length)
{
    if(length == 1) return 0;
    size_t length_log2 = LOG2(length);
    if(length_log2 > map->orders) length_log2 = map->orders - 1;
    return length_log2;
}

void
buddy_mark_range(struct BuddyMap *map, size_t start, size_t end)
{
    size_t l = s_buddy_layer_bestfit(map, end - start);
    s_buddy_op_range(map, l, start, end, buddy_bit_mark);
}
void 
buddy_free_range(struct BuddyMap *map, size_t start, size_t end)
{ 
    size_t l = s_buddy_layer_bestfit(map, end - start);
    s_buddy_op_range(map, l, start, end, buddy_bit_free);
}

static intmax_t
s_buddy_top_firstfree(struct BuddyMap *map)
{
    size_t layer = map->orders - 1;
    size_t layer_bits = map->bits >> layer;
    size_t layer_blocks = layer_bits / BUDDY_BLOCK_BITS;
    for(size_t i = 0; i < layer_blocks; i++) {
        if(map->blocks[layer][i] == (uintptr_t)-1) continue;
        for(size_t j = 0; j < BUDDY_BLOCK_BITS; j++) {
            size_t bit = (i * BUDDY_BLOCK_BITS) + j;
            if(buddy_bit_test(map, layer, bit)) continue;
            return bit;
        }
    }
    return -1;
}

static intmax_t
s_buddy_top_alloc_contig(struct BuddyMap *map, size_t length)
{
    size_t layer = map->orders - 1;
    size_t layer_bits = map->bits >> layer;

    size_t contig_base = 0;
    size_t contig_bits = 0;
    for(size_t i = 0; i < layer_bits; i++) {
        if(buddy_bit_test(map, layer, i)) {
            contig_base = i + 1;
            contig_bits = 0;
            continue;
        }
        if(++contig_bits == length) {
            return contig_base;
        }
    }
    return -1;
}

intmax_t
buddy_alloc(struct BuddyMap *map, size_t length)
{
    size_t layer = s_buddy_layer_bestfit(map, length);
    size_t layerw = 1ULL << layer;
    size_t allocw = length / layerw;

    if(allocw == 0) allocw = 1;
    if(allocw > 1)
        return s_buddy_top_alloc_contig(map, length);

    intmax_t alloci = s_buddy_top_firstfree(map);
    if(alloci < 0) {
        klogf("Top layer is full!\n");
        return -1;
    }

    for(int wlayer = map->orders - 2; wlayer > layer; wlayer--) {
        alloci <<= 1;
        if(buddy_bit_test(map, wlayer, alloci)) {
            alloci++;
        }
    }
    alloci <<= 1;
    if(buddy_bit_test(map, layer, alloci)) alloci++;

    s_buddy_op_range(map, layer, alloci, alloci + length, buddy_bit_mark);
    return alloci;
}

void
buddy_free(struct BuddyMap *map, size_t i, size_t length)
{
    buddy_free_range(map, i, i + length);
}