#include "sw_mlkem_alloc.h"

#include <stdint.h>

#define SW_MLKEM_ALLOC_ALIGN 32u
#define SW_MLKEM_ALLOC_SIZE  (64u * 1024u)

static uint8_t sw_mlkem_arena[SW_MLKEM_ALLOC_SIZE] __attribute__((aligned(SW_MLKEM_ALLOC_ALIGN)));
static size_t sw_mlkem_arena_pos;
static unsigned sw_mlkem_arena_active;

static size_t align_up(size_t n)
{
    return (n + (SW_MLKEM_ALLOC_ALIGN - 1u)) & ~(SW_MLKEM_ALLOC_ALIGN - 1u);
}

void *sw_mlkem_alloc(size_t size)
{
    size_t pos = align_up(sw_mlkem_arena_pos);
    size_t need = align_up(size);

    if (pos + need > SW_MLKEM_ALLOC_SIZE)
        return 0;

    sw_mlkem_arena_pos = pos + need;
    sw_mlkem_arena_active++;
    return &sw_mlkem_arena[pos];
}

void sw_mlkem_free(void *ptr, size_t size)
{
    (void)ptr;
    (void)size;

    if (sw_mlkem_arena_active != 0)
        sw_mlkem_arena_active--;

    if (sw_mlkem_arena_active == 0)
        sw_mlkem_arena_pos = 0;
}