#include "arena.h"

#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>

#define ARENA_DEFAULT_CAPACITY ((size_t)4096)

struct ArenaBlock {
    ArenaBlock *next;
    size_t used;
    size_t capacity;
    max_align_t alignment;
    uint8_t data[];
};

void arena_init(Arena *arena) {
    arena->head = NULL;
}

void arena_destroy(Arena *arena) {
    ArenaBlock *block = arena->head;
    while (block != NULL) {
        ArenaBlock *next = block->next;
        free(block);
        block = next;
    }
    arena->head = NULL;
}

void *arena_allocate(Arena *arena, size_t count, size_t item_size) {
    const size_t alignment = _Alignof(max_align_t);
    size_t bytes;
    size_t offset = 0;
    ArenaBlock *block = arena->head;

    if (count == 0 || item_size == 0) {
        return NULL;
    }
    if (count > SIZE_MAX / item_size) {
        return NULL;
    }
    bytes = count * item_size;
    if (block != NULL) {
        if (block->used > SIZE_MAX - (alignment - 1U)) {
            return NULL;
        }
        offset = block->used +
                 (alignment - block->used % alignment) % alignment;
    }
    if (block == NULL || offset > block->capacity ||
        bytes > block->capacity - offset) {
        size_t capacity = bytes > ARENA_DEFAULT_CAPACITY
                              ? bytes : ARENA_DEFAULT_CAPACITY;
        if (capacity > SIZE_MAX - offsetof(ArenaBlock, data)) {
            return NULL;
        }
        block = malloc(offsetof(ArenaBlock, data) + capacity);
        if (block == NULL) {
            return NULL;
        }
        block->next = arena->head;
        block->used = 0;
        block->capacity = capacity;
        arena->head = block;
        offset = 0;
    }
    block->used = offset + bytes;
    return block->data + offset;
}