#include <assert.h>
#include <string.h>
#include "lpool.h"

#define ALIGN_UP(n, alignment) (((n) + (__typeof__(n))(alignment) - 1) & ~((__typeof__(n))(alignment) - 1))
#define NULL_BLOCK (0xffffffff)

static inline LPoolBlock *lpool_next(LPool *pool, LPoolBlock *block)
{
    if (block->next == NULL_BLOCK) {
        return NULL;
    }
    return (LPoolBlock *)&pool->start[block->next];
}

static inline uint32_t lpool_offset(LPool *pool, LPoolBlock *block)
{
    if (block == NULL) {
        return NULL_BLOCK;
    }
    return (uint8_t *)block - pool->start;
}

void lpool_init(LPool *pool, void *buffer, size_t buffer_size, uint32_t block_size, uint32_t align)
{
    assert(align > 0);
    assert((align & (align - 1)) == 0);  // power of 2

    // enforce minimum for storing free list offsets

    if (align < _Alignof(LPoolBlock)) {
        align = _Alignof(LPoolBlock);
    }

    if (block_size < sizeof(LPoolBlock)) {
        block_size = sizeof(LPoolBlock);
    }
    // ensure initial alignment
    block_size = ALIGN_UP(block_size, align);

    // align first usable block
    uintptr_t buffer_addr = (uintptr_t)buffer;
    uintptr_t first_addr = ALIGN_UP(buffer_addr, align);
    size_t padding = first_addr - buffer_addr;

    if (padding + block_size > buffer_size) {
        pool->capacity = 0;
        pool->free_list = NULL;
    } else {
        size_t usable = buffer_size - padding;
        pool->capacity = usable / block_size;

        // build intrusive free list
        LPoolBlock *block = (LPoolBlock *)first_addr;
        pool->free_list = block;
        for (uint32_t i = 1; i < pool->capacity; i++) {
            uint32_t next = i * block_size;
            block->next = next;
            block = (LPoolBlock *)(first_addr + next);
        }
        block->next = NULL_BLOCK;  // last element
    }
    pool->block_size = block_size;
    pool->align = align;
    pool->start = buffer;
    pool->used = 0;
}

void *lpool_alloc(LPool *pool)
{
    if (pool->free_list == NULL) {
        return NULL;
    }
    LPoolBlock *block = pool->free_list;
    LPoolBlock *next = lpool_next(pool, block);
    pool->free_list = next;

#ifdef DEBUG
    memset(block, 0xCB, pool->block_size);
#endif
    pool->used++;
    return block->data;
}

void lpool_free(LPool *pool, void *ptr)
{
    if (ptr == NULL) {
        return;
    }
#ifdef DEBUG
    assert(pool->used > 0);
    uint8_t *end = pool->start + (pool->capacity * pool->block_size);
    assert((uint8_t *)ptr >= pool->start && (uint8_t *)ptr < end);
    assert((((uintptr_t)ptr - (uintptr_t)pool->start) & (pool->block_size - 1)) == 0);

    // double-free detection
    for (LPoolBlock *block = pool->free_list; block != NULL; block = lpool_next(pool, block)) {
        if (block == ptr) {
            assert(0 && "Double free detected in pool!");
            return;
        }
    }
    memset(ptr, 0xDD, pool->block_size);
#endif
    LPoolBlock *block = ptr;
    // When not in use, store pointer to next free element inline in the block
    block->next = lpool_offset(pool, pool->free_list);
    pool->free_list = block;
    pool->used--;
}

void lpool_reset(LPool *pool)
{
    if (pool->capacity == 0) {
        return;
    }

    uintptr_t first_addr = (uintptr_t)pool->start;

    // build intrusive free list
    LPoolBlock *block = (LPoolBlock *)first_addr;
    pool->free_list = block;
    for (size_t i = 1; i < pool->capacity; ++i) {
        uint32_t next = i * pool->block_size;
        block->next = next;
        block = (LPoolBlock *)(first_addr + next);
    }

    block->next = NULL_BLOCK;  // last element
    pool->used = 0;
}