#ifndef UC_SARENA_H_
#define UC_SARENA_H_

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

/** @file
 * Static arena allocator.
 *
 *  An UCSArena allocates memory from a provided buffer.
 *
 * The backing buffer should be suitable aligned to the types
 * that will be allocated from it. Use .e.g
 * @code
 * _Alignas(16) char mem[1024];
 * @endcode
 * as the backing buffer.
 *
 * Once allocated, a piece of memory obtained from the
 * allocator cannot be free'd individually, only the entire SArena can
 * be free'd/reset
 */
typedef struct UCSArena UCSArena;
struct UCSArena {
    unsigned char *start;  // start of user buffer
    unsigned char *curr;   // next alloc point
    unsigned char *end;    // one past end of user buffer
    int tmp_allocs;        // no of uc_sar_tmp_begin calls
};

typedef struct UCSTmpAlloc UCSTmpAlloc;
struct UCSTmpAlloc {
    unsigned char *alloc_point;
    UCSArena *arena;
};
/** Initialize a UCSArena with predefined backing array.
 *
 * @data must be an array type
 * Use as
 * @code
 * _Alignas(16) char v[128];
 * UCSArena a = UC_SAR_STATIC_INITIALIZER(v);
 * @endcode
 */
#define UC_SAR_STATIC_INITIALIZER(data) UC_SAR_INITIALIZER((data), sizeof(data))

/** Initializer for UCSArena with predefined memory
 * Use as
 * @code
 * char v[128];
 * UCSArena a = UC_SAR_INITIALIZER(v, sizeof v);
 * @endcode
 * or
 * @code
 * char *m = malloc(128);
 * UCSArena a = UC_SAR_INITIALIZER(m, 128);
 * @endcode
 */
#define UC_SAR_INITIALIZER(data, sz) {\
    .start = (data),\
    .curr = (data),\
    .end=(data) + sz,\
    .tmp_allocs = 0\
}

/** Initialize a UCSArena with an array as the backing memory to allocate from.
 * The macro can be used at local scope to allocate memory automatic on the stack,
 * or at global scope where the backing array will have static storage duration
 *
 * Use as:
 * @code
 * UC_SAR_ALLOC_DEF(my_allocator, 1024)
 * struct foo *foo = uc_sar_alloc(&my_allocator);
 * if (foo == NULL) {
 *     // out of memory
 * }
 * @endcode
 *
 * @note, the backing buffer has alignment of 16. Use manual init if a different alignment
 * is needed.
 * @param var_name Variable name for the allocator
 * @param size_  byte size of the backing array
 */
#define UC_SAR_ALLOC_DEF(var_name, size_) \
    struct {                              \
        _Alignas(16) char memory[size_];  \
    } var_name##_arena_memory;            \
    UCSArena var_name = UC_SAR_STATIC_INITIALIZER(var_name##_arena_memory.memory)

/** Allocate space for @type_ , returned memory is aligned to requirements
 * of @type.
 *
 * Use as:
 * @code
 * UC_SAR_ALLOC_DEF(my_allocator, 1024)
 * struct foo *foo = UC_SAR_ALLOC_ALIGNED(&my_allocator, struct foo);
 * // Array types can be used too:
 * struct Foo *str = UC_SAR_ALLOC_ALIGNED(&my_allocator, struct Foo[32]);
 * @endcode
 */
#define UC_SAR_ALLOC_ALIGNED(allocator, type_) uc_sar_alloc_aligned((allocator), sizeof(type_), _Alignof(type_))

// Initialize an UCSArena to allocate from @mem which is @sz bytes
static inline void uc_sar_init(UCSArena *arena, void *memory, size_t sz)
{
    arena->start = arena->curr = (unsigned char *)memory;
    arena->end = arena->start + sz;
}

// begin temporary allocation. All allocations afterwards will be released in uc_sar_tmp_end
static inline UCSTmpAlloc uc_sar_tmp_begin(UCSArena *arena)
{
    UCSTmpAlloc state = {.alloc_point = arena->curr, .arena = arena};
    arena->tmp_allocs++;
    return state;
}

// end temporary allocations, releasing all allocations since the most recent uc_sar_tmp_begin
static inline void uc_sar_tmp_end(const UCSTmpAlloc *state)
{
    UCSArena *arena = state->arena;
    assert(arena->tmp_allocs > 0);
    arena->curr = state->alloc_point;
    arena->tmp_allocs--;
}

// Allocate @sz bytes from the UCSArena. Returns NULL if not enough space
static inline void *uc_sar_alloc(UCSArena *arena, size_t sz)
{
    if (sz > (arena->curr - arena->end)) {
        return NULL;
    }
    void *start = arena->curr;
    arena->curr += sz;
    return start;
}

// Allocate @sz bytes from the UCSArena, returned pointer is aligned to @align
// Returns NULL if not enough space
static inline void *uc_sar_alloc_aligned(UCSArena *arena, size_t sz, unsigned int align) [[gnu::alloc_align(3)]]
{
    assert(align != 0);
    assert((align & (align - 1)) == 0);  // power of 2

    uintptr_t curr = (uintptr_t)arena->curr;
    uintptr_t aligned = (curr + (align - (uintptr_t)1)) & ~(align - (uintptr_t)1);
    unsigned char *start = (unsigned char *)aligned;

    if (sz > (arena->curr - arena->end)) {

        return NULL;
    }

    arena->curr = start + sz;

    return start;
}

// reset arena, making all memory available for allocation
// also invalidates any UCSTmpAlloc instances.
static inline void uc_sar_reset(UCSArena *arena)
{
    arena->curr = arena->start;
}

// Arena size in bytes
static inline size_t uc_sar_size(const UCSArena *arena)
{
    return (size_t)(arena->end - arena->start);
}

// Used bytes of the arena
static inline size_t uc_sar_used(const UCSArena *arena)
{
    return (size_t)(arena->curr - arena->start);
}

// Available unused bytes in the arena
static inline size_t uc_sar_available(const UCSArena *arena)
{
    return uc_sar_size(arena) - uc_sar_used(arena);
}

#endif