#include "kpl.h"

void kpl_memory_allocator_init(kpl_memory_allocator *memory_allocator) {
    for (int32_t slab_pow2 = KPL_MEMORY_ALLOCATOR_MIN_SIZE_POW2;
        slab_pow2 <= KPL_MEMORY_ALLOCATOR_MAX_SIZE_POW2; slab_pow2++) {
        kpl_memory_allocator_slab_group *slab_group = &memory_allocator->slabs[slab_pow2];
        slab_group->slab_size_pow2 = slab_pow2;
        slab_group->slab_head = slab_group->slab_tail = NULL;
        kpl_object_queue_init(&slab_group->object_pool);
    }
}

static void kpl_memory_allocator_slab_free(kpl_memory_allocator_slab_group *slab_group) {
    kpl_memory_allocator_slab *slab_head = slab_group->slab_head;
    while (slab_head) {
        kpl_memory_allocator_slab *slab_tmp = slab_head;
        slab_head = slab_head->slab_next;
        munmap(slab_tmp->buffer, slab_tmp->buffer_size);
        free(slab_tmp);
    }
}

void kpl_memory_allocator_free(kpl_memory_allocator *memory_allocator) {
    for (int32_t slab_pow2 = KPL_MEMORY_ALLOCATOR_MIN_SIZE_POW2;
        slab_pow2 < KPL_MEMORY_ALLOCATOR_SLAB_COUNT; slab_pow2++) {
        kpl_memory_allocator_slab_group *slab_group = &memory_allocator->slabs[slab_pow2];
#ifdef KPL_VALGRIND
        kpl_object *object = NULL;
        while ((object = kpl_object_queue_dequeue(&slab_group->object_pool)))
            VALGRIND_FREELIKE_BLOCK(object, 0);
#endif
        kpl_memory_allocator_slab_free(slab_group);
    }
}

static int32_t kpl_memory_allocator_get_size_pow2(int32_t size) {
    return size > 1 ? 32 - __builtin_clz(size - 1) : 0;
}

static const int32_t pages_per_pow2[] = {
    [5]     = 5,
    [6]     = 10,
    [7]     = 15,
    [8]     = 20,
    [9]     = 25,
    [10]    = 30,
    [11]    = 35,
    [12]    = 40,
    [13]    = 2,
    [14]    = 4,
    [15]    = 8,
    [16]    = 16,
    [17]    = 32,
    [18]    = 64,
    [19]    = 128,
    [20]    = 256,
    [21]    = 512,
    [22]    = 1024
};

static kpl_memory_allocator_slab *kpl_memory_allocator_slab_init(int32_t size_pow2) {
    kpl_memory_allocator_slab *slab = calloc(1, sizeof(kpl_memory_allocator_slab));
    slab->buffer_size = getpagesize() * pages_per_pow2[size_pow2];
    slab->buffer = mmap(NULL, slab->buffer_size, PROT_READ | PROT_WRITE, MAP_ANONYMOUS | MAP_SHARED, -1, 0);
    return slab;
}

static void *kpl_memory_allocator_slab_bump(kpl_memory_allocator_slab *slab, int32_t size) {
    void *object = slab->buffer + slab->buffer_used;
    slab->buffer_used += size;
    return object;
}

static kpl_object *kpl_memory_allocator_slab_allocate(kpl_memory_allocator_slab_group *slab_group,
    int32_t thread_id) {
    kpl_memory_allocator_slab *slab = slab_group->slab_tail;
    int32_t size = 1 << slab_group->slab_size_pow2;
    if (!slab || slab->buffer_used + size > slab->buffer_size) {
        slab = kpl_memory_allocator_slab_init(slab_group->slab_size_pow2);
        if (!slab_group->slab_head)
            slab_group->slab_head = slab_group->slab_tail = slab;
        else
            slab_group->slab_tail = slab_group->slab_tail->slab_next = slab;
    }
    kpl_object *object = kpl_memory_allocator_slab_bump(slab, size);
#ifdef KPL_VALGRIND
    VALGRIND_MALLOCLIKE_BLOCK(object, size, 0, 0);
#endif
    __atomic_store_n(&object->queue_next, NULL, __ATOMIC_RELAXED);
    object->object_size = size;
    object->thread_id = thread_id;
    return object;
}

void *kpl_memory_allocator_allocate(kpl_memory_allocator *memory_allocator, int32_t size, int32_t thread_id) {
    if (size < 1 << KPL_MEMORY_ALLOCATOR_MIN_ROUND_SIZE_POW2)
        return NULL;
    int32_t size_pow2 = kpl_memory_allocator_get_size_pow2(size);
    if (size_pow2 < KPL_MEMORY_ALLOCATOR_MIN_SIZE_POW2)
        size_pow2 = KPL_MEMORY_ALLOCATOR_MIN_SIZE_POW2;
    kpl_memory_allocator_slab_group *slab_group = &memory_allocator->slabs[size_pow2];
    kpl_object *object = kpl_object_queue_dequeue(&slab_group->object_pool);
    if (!object)
        object = kpl_memory_allocator_slab_allocate(slab_group, thread_id);
    return object;
}

void kpl_memory_allocator_return(kpl_memory_allocator *memory_allocator, void *object) {
    int32_t size_pow2 = kpl_memory_allocator_get_size_pow2(((kpl_object*) object)->object_size);
    kpl_object_queue_enqueue(&memory_allocator->slabs[size_pow2].object_pool, object);
}