// arena.h - v1.2.0 - MIT License
// single header library for region-based memory management.
//
// [License and changelog]
//
//     See end of file.
//
// [Single header library usage]
//
//     Including this header file as such does what you normally expect.
//
//     If you also need to include all of the function implementations you
//     need to #define ARENA_IMPLEMENTATION before including the header file
//     like this:
//
//         #define ARENA_IMPLEMENTATION
//         #include "arena.h"
//
// [Compile-time options]
//
//     Note: every compile-time option listed here should be configured before
//           the #include "arena.h"
//
//     #define ARENA_REGION_CAPACITY new_region_capacity_in_bytes (4096)
//
//         This macro defines the default capacity for arena regions.
//
//     #define ARENA_ALIGNMENT alignment (8)
//
//         This macro defines the memory alignment used for allocations.
//
//     #define ARENA_NO_ASSERT
//
//         This macro disables assertions that happen when allocations fail.
//         The disabled assertions will be replaced with a simple condition
//         that returns NULL when it fails. This macro affects every function
//         that returns a pointer.
//
//     #define ARENA_ASSERT my_assert
//
//         This macro defines an alternative function for assertions. This has
//         no effect when ARENA_NO_ASSERT is defined.
//
//     #define ARENA_MALLOC my_malloc
//     #define ARENA_FREE my_free
//
//         These macros define alternative functions for dynamic allocation
//         and deallocation. They are only used for managing Arena_Region
//         structures. If redefined ensure that your functions consider memory
//         alignment.
//
// [Function documentation]
//
// Note: every function that starts with "arena__" is considered a private
// function of this library.
//
// In this library, arenas are implemented as linked lists of regions. Each
// region will contain the allocated buffers. The following are all of the
// functions
//
//     Arena arena_create(size_t region_capacity)
//
// This function is optional (see [Example] section): initializes an arena
// with a specified region capacity. It will override the default
// ARENA_REGION_CAPACITY.
//
//     void arena_free(Arena *a)
//
// This function frees all of the regions allocated in the specified arena
// which invalidates all the pointers associated with the arena. The arena can
// still be reused. If the arena was created with arena_create() the
// configured region_capacity will be retained.
//
//     void arena_reset(Arena *a)
//
// This function resets an arena by keeping all of the regions allocated but
// invalidates all of the pointers associated with the specified arena.
// Warning: this functions may cause fragmentation, consider setting an
// appropriate region capacity.
//
//     void *arena_alloc(Arena *a, size_t size)
//
// This function allocates a buffer of <size> bytes into the specified arena
// and returns it's pointer.
//
//     void *arena_copy(Arena *a, const void *src, size_t size);
//
// This function allocates a buffer of <size> bytes into the arena and copies
// the contents of <src>. Returns the pointer of the new copy.
//
//     char *arena_strdup(Arena *a, const char *s);
//
// This functions acts like the strdup() function but it instead allocates on
// the specified arena. It duplicates a string in the arena.
//
//     char *arena_sprintf(Arena *a, const char *fmt, ...)
//
// This function internally uses vsnprintf() to store the resulting formatted
// string in the specified arena. Returns NULL on error.
//
//     Arena_Marker arena_snapshot(const Arena *a)
//
// This function creates a snapshot of the current state of the arena and
// returns it. This is useful when you only need to allocate memory on the arena
// temporarily. To restore the returned state, see arena_rewind().
//
//     void arena_rewind(Arena *a, Arena_Marker marker)
//
// This function restores the arena to the specified state, invalidating all
// previously allocated buffers since the creation of the snapshot.
//
// [Notes]
//
// This library is not Thread-safe.
//
// [Example]
//
#if 0
#define ARENA_IMPLEMENTATION
#include "arena.h"

int main(void)
{
    Arena a = {0}; // or arena_create(...) to specify a custom region
                   // capacity just for this arena

    int *x = arena_alloc(&a, sizeof(*x) * 32);
    float *y = arena_alloc(&a, sizeof(*y) * 512);

    arena_free(&a); // The arena can still be reused
    return 0;
}
#endif // [Example]

#ifndef ARENA_H_
#define ARENA_H_

#include <stdalign.h>
#include <stdarg.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>

#ifndef ARENA_REGION_CAPACITY
#    define ARENA_REGION_CAPACITY (8*1024)
#endif // ARENA_REGION_CAPACITY

#ifndef ARENA_ALIGNMENT
#    define ARENA_ALIGNMENT alignof(max_align_t)
#endif // ARENA_ALIGNMENT

#ifndef ARENA_NO_ASSERT
#    ifndef ARENA_ASSERT
#        include <assert.h>
#        define ARENA_ASSERT assert
#    endif // ARENA_ASSERT
#else
#    define ARENA_ASSERT(...) ((void)0)
#endif // ARENA_NO_ASSERT

#ifndef ARENA_MALLOC
#    include <stdlib.h>
#    define ARENA_MALLOC malloc
#endif // ARENA_MALLOC

#ifndef ARENA_FREE
#    include <stdlib.h>
#    define ARENA_FREE free
#endif // ARENA_FREE

#ifdef __cplusplus
extern "C" { // Prevent name mangling of functions
#endif // __cplusplus

typedef struct Arena_Region Arena_Region;

struct Arena_Region {
    Arena_Region *next;
    size_t count;
    size_t capacity;
    alignas(ARENA_ALIGNMENT) uint8_t data[];
};

typedef struct {
    Arena_Region *region;
    size_t count;
} Arena_Marker;

typedef struct Arena {
    Arena_Region *head;
    Arena_Region *tail;
    size_t region_capacity;
} Arena;

// Arena management
Arena arena_create(size_t region_capacity);
void arena_free(Arena *a);
void arena_reset(Arena *a);

// Allocation and other utilities
void *arena_alloc(Arena *a, size_t size);
void *arena_copy(Arena *a, const void *src, size_t size);
char *arena_strdup(Arena *a, const char *s);
char *arena_sprintf(Arena *a, const char *fmt, ...);

// Snapshot management
Arena_Marker arena_snapshot(const Arena *a);
void arena_rewind(Arena *a, Arena_Marker marker);

// Private functions
size_t arena__align(size_t n);

#ifdef __cplusplus
}
#endif // __cplusplus

#endif // ARENA_H_

#ifdef ARENA_IMPLEMENTATION

#ifdef __cplusplus
extern "C" { // Prevent name mangling of functions
#endif // __cplusplus

Arena arena_create(size_t region_capacity)
{
    Arena a;
    memset(&a, 0, sizeof(a));
    a.region_capacity = region_capacity;
    return a;
}

void arena_free(Arena *a)
{
    if (a == NULL) {
        return;
    }

    Arena_Region *cur = a->head;
    while (cur != NULL) {
        Arena_Region *next = cur->next;
        ARENA_FREE(cur);
        cur = next;
    }

    a->head = NULL;
    a->tail = NULL;
    // a->region_capacity = 0;
}

void arena_reset(Arena *a)
{
    if (a == NULL) {
        return;
    }

    for (Arena_Region *cur = a->head; cur != NULL; cur = cur->next) {
        cur->count = 0;
    }

    a->tail = a->head;
}

void *arena_alloc(Arena *a, size_t size)
{
    if (a == NULL || size == 0) {
        ARENA_ASSERT(false && "Invalid parameters");
        return NULL;
    }

    // Ensure memory alignment of requested size
    size = arena__align(size);

    // Check for overflow
    if (size > SIZE_MAX - sizeof(Arena_Region)) {
        ARENA_ASSERT(false && "Requested size too large");
        return NULL;
    }

    size_t region_capacity = (a->region_capacity == 0
                                ? ARENA_REGION_CAPACITY : a->region_capacity);

    // Calculate and check region allocation size
    size_t alloc_size = (size > region_capacity ? size : region_capacity);
    if (alloc_size > SIZE_MAX - sizeof(Arena_Region)) {
        ARENA_ASSERT(false && "Region size overflow");
        return NULL;
    }

    // Empty arena: allocate first region
    if (a->head == NULL) {
        a->head = (Arena_Region*)ARENA_MALLOC(sizeof(*a->head) + alloc_size);
        if (a->head == NULL) {
            ARENA_ASSERT(false && "Allocation failed");
            return NULL;
        }

        a->head->next = NULL;
        a->head->count = size;
        a->head->capacity = alloc_size;
        a->tail = a->head;
        return a->head->data;
    }

    // Find first suitable region
    Arena_Region *r = a->tail;
    while (r != NULL && size > r->capacity - r->count) {
        r = r->next;
    }

    // If found allocate on it
    if (r != NULL) {
        if (r->count > SIZE_MAX - size) {
            ARENA_ASSERT(false && "Region count overflow");
            return NULL;
        }
        r->count += size;
        return r->data + r->count - size;
    }

    // If not found append a new region
    a->tail->next = (Arena_Region*)ARENA_MALLOC(sizeof(*a->tail) + alloc_size);
    if (a->tail->next == NULL) {
        ARENA_ASSERT(false && "Allocation failed");
        return NULL;
    }

    a->tail = a->tail->next;
    a->tail->next = NULL;
    a->tail->count = size;
    a->tail->capacity = alloc_size;
    return a->tail->data;
}

void *arena_copy(Arena *a, const void *src, size_t size)
{
    if (a == NULL || src == NULL) {
        ARENA_ASSERT(false && "Invalid parameters");
        return NULL;
    }

    void *copy = arena_alloc(a, size);
    if (copy == NULL) {
        return NULL;
    }

    memcpy(copy, src, size);

    return copy;
}

char *arena_strdup(Arena *a, const char *s)
{
    if (a == NULL || s == NULL) {
        ARENA_ASSERT(false && "Invalid parameters");
        return NULL;
    }

    return (char*)arena_copy(a, s, strlen(s) + 1);
}

char *arena_sprintf(Arena *a, const char *fmt, ...)
{
    if (a == NULL || fmt == NULL) {
        ARENA_ASSERT(false && "Invalid parameters");
        return NULL;
    }

    va_list args;
    va_start(args, fmt);

    Arena_Region *r = a->tail;
    char *result = NULL;
    size_t available = 0;

    if (r != NULL) {
        available = r->capacity - r->count;
        if (available > 0) {
            result = (char*)r->data + r->count;
        }
    }

    // Check if result fits inside tail's region available space
    va_list args_copy;
    va_copy(args_copy, args);
    int n = vsnprintf(result, available, fmt, args_copy);
    va_end(args_copy);

    if (n < 0) {
        va_end(args);
        return NULL;
    }

    size_t needed = (size_t)n + 1;

    // If it fits commit and return
    if (needed <= available) {
        r->count += arena__align(needed);
        va_end(args);
        return result;
    }

    // Fallback: if it didn't fit, allocate and printf
    result = (char*)arena_alloc(a, needed);
    if (result == NULL) {
        va_end(args);
        return NULL;
    }

    n = vsnprintf(result, needed, fmt, args);
    va_end(args);

    if (n < 0) {
        return NULL;
    }

    return result;
}

Arena_Marker arena_snapshot(const Arena *a)
{
    Arena_Marker marker = {0};

    if (a == NULL) {
        ARENA_ASSERT(false && "Invalid parameters");
        return marker;
    }

    marker.region = a->tail;
    marker.count = (a->tail != NULL ? a->tail->count : 0);

    return marker;
}

void arena_rewind(Arena *a, Arena_Marker marker)
{
    if (a == NULL) {
        ARENA_ASSERT(false && "Invalid parameters");
        return;
    }

    // If arena was empty reset it
    if (marker.region == NULL) {
        arena_reset(a);
        return;
    }

    a->tail = marker.region;
    a->tail->count = marker.count;

    // Clear new regions in case they were appended
    for (Arena_Region *cur = a->tail->next; cur != NULL; cur = cur->next) {
        cur->count = 0;
    }
}

size_t arena__align(size_t n)
{
    return (n + (ARENA_ALIGNMENT - 1)) & ~(ARENA_ALIGNMENT - 1);
}

#ifdef __cplusplus
}
#endif // __cplusplus

#endif // ARENA_IMPLEMENTATION

/*
 * Revision history:
 *
 *     1.2.0 (2026-04-08) New helper functions: arena_sprintf(),
 *                        arena_snapshot() and arena_rewind()
 *     1.1.3 (2026-01-14) Align Arena_Region data flexible array member
 *     1.1.2 (2026-01-10) Minor changes; check result of alloc in arena_copy()
 *     1.1.1 (2025-11-18) Implement memory alignment; improve docs
 *     1.1.0 (2025-11-18) New helper functions: arena_copy(), arena_strdup()
 *     1.0.3 (2025-09-10) Renamed parameter bytes to size in arena_alloc()
 *     1.0.2 (2025-09-06) Bug fixes; new ARENA_NO_ASSERT macro
 *     1.0.1 (2025-08-02) Prevent name mangling of functions; don't reset
 *                        region_capacity in arena_free()
 *     1.0.0 (2025-07-24) Initial release
 */

/*
 * MIT License
 *
 * Copyright (c) 2025 seajee
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */