#include <stdint.h>
#include <stddef.h>
#include <assert.h>
#include "lalloc.h"

#define ALIGNMENT sizeof(BlockHeader)
#define ALIGN_UP(n) (((n) + (__typeof__(n))(ALIGNMENT) - 1) & ~((__typeof__(n))(ALIGNMENT) - 1))
#define ALIGN_DOWN(n) ((n) & ~((__typeof__(n))(ALIGNMENT) - 1))

#define NULL_BLOCK (0xffffffff)
#define BLOCK_FREE ((uint32_t)1)

static inline void block_set_free(BlockHeader *block)
{
    block->size_and_flags |= BLOCK_FREE;
}

static inline void block_set_used(BlockHeader *block)
{
    block->size_and_flags &= ~BLOCK_FREE;
}

static inline void block_set_size(BlockHeader *block, uint32_t size)
{
    uint32_t flags = block->size_and_flags & BLOCK_FREE;
    block->size_and_flags = size | flags;
}

uint32_t block_size(const BlockHeader *block)
{
    return block->size_and_flags & ~BLOCK_FREE;
}

int block_is_free(const BlockHeader *block)
{
    return (block->size_and_flags & BLOCK_FREE) != 0;
}

inline BlockHeader *block_from_offset(const LAlloc *allocator, uint32_t offset)
{
    if (offset == NULL_BLOCK) {
        return NULL;
    }

    return (BlockHeader *)((uintptr_t)allocator->base + offset);
}

static inline uint32_t bh_offset(LAlloc *allocator, BlockHeader *block)
{
    if (block == NULL) {
        return NULL_BLOCK;
    }

    return (uint32_t)((uintptr_t)block - (uintptr_t)allocator->base);
}

static void maybe_split_block(LAlloc *allocator, BlockHeader *block, uint32_t size)
{
    // assumes size is already aligned
    uint32_t bsize = block_size(block);

    if (bsize < size + sizeof(BlockHeader) + ALIGNMENT) {
        return;
    }

    uint8_t *base = (uint8_t *)block;
    BlockHeader *new_block = (BlockHeader *)(base + sizeof(BlockHeader) + size);

    uint32_t new_size = bsize - size - sizeof(BlockHeader);

    block_set_size(block, size);
    block_set_used(block);

    new_block->size_and_flags = new_size | BLOCK_FREE;
    new_block->next = block->next;
    new_block->prev = bh_offset(allocator, block);

    BlockHeader *next = block_from_offset(allocator, new_block->next);
    if (next) {
        next->prev = bh_offset(allocator, new_block);
    }

    block->next = bh_offset(allocator, new_block);
}

static void merge_with_next(LAlloc *allocator, BlockHeader *block)
{
    BlockHeader *next = block_from_offset(allocator, block->next);
    if (next == NULL || !block_is_free(next)) return;

    uint32_t new_size = block_size(block) + sizeof(BlockHeader) + block_size(next);

    block_set_size(block, new_size);
    block->next = next->next;
    BlockHeader *nextnext = block_from_offset(allocator, next->next);
    if (nextnext) nextnext->prev = bh_offset(allocator, block);
    // note - caller have to adjust next_alloc on return
}

void lalloc_init(LAlloc *allocator, void *backing_store, uint32_t size)
{
    // backing store must be aligned, otherwise  all alignment guarantees are lies
    assert(((uintptr_t)backing_store & (ALIGNMENT - 1)) == 0);
    assert(size >= sizeof(BlockHeader));
    assert(size < 0xfffffff0);

    allocator->base = backing_store;
    allocator->size = size;

    BlockHeader *first = backing_store;
    // allocator logic requires the first block to be properly aligned
    // as well as size being eligned.
    uint32_t usable = ALIGN_DOWN(size) - sizeof(BlockHeader);
    block_set_free(first);
    block_set_size(first, usable);
    first->next = NULL_BLOCK;
    first->prev = NULL_BLOCK;

    allocator->heap_start = first;
    allocator->next_alloc = first;
}

void *lalloc(LAlloc *allocator, uint32_t size)
{
    if (size > UINT32_MAX - (ALIGNMENT - 1)) {
        return NULL;
    }
    size = ALIGN_UP(size);

    BlockHeader *start = allocator->next_alloc ? allocator->next_alloc : allocator->heap_start;
    BlockHeader *curr = start;

    do {
        if (block_is_free(curr) && block_size(curr) >= size) {
            maybe_split_block(allocator, curr, size);
            block_set_used(curr);
            BlockHeader *next = block_from_offset(allocator, curr->next);
            allocator->next_alloc = next ? next : allocator->heap_start;
            return (uint8_t *)curr + sizeof(BlockHeader);
        }
        BlockHeader *next = block_from_offset(allocator, curr->next);
        curr = next ? next : allocator->heap_start;
    } while (curr != start);

    return NULL;
}

void lfree(LAlloc *allocator, void *ptr)
{
    if (!ptr) {
        return;
    }

    BlockHeader *block = (BlockHeader *)((uint8_t *)ptr - sizeof(BlockHeader));

    assert((uint8_t *)ptr - sizeof(BlockHeader) >= allocator->base && (uint8_t *)ptr < allocator->base + allocator->size);
    // double free detection
    assert(!block_is_free(block));

    block_set_free(block);
    merge_with_next(allocator, block);

    BlockHeader *prev = block_from_offset(allocator, block->prev);
    if (prev != NULL && block_is_free(prev)) {
        merge_with_next(allocator, prev);
        allocator->next_alloc = prev;
    } else {
        allocator->next_alloc = block;
    }
}