// ufn_arena.h: An arena allocator with maximum usage limiter.
// Under MIT-License. Copyright by Eric-Paul Ickhorn (epickh).

#ifndef UFN_ARENA_H
#define UFN_ARENA_H

#include <stdint.h>

#define UFN_UNLIMITED_ARENA 0xffffffff

typedef struct ufn_arena ufn_arena_s;

/// @brief Arena structure which is being acted upon by the arena
///        functions. Manual modification may lead to undesired behaviour.
struct ufn_arena
{
    uint32_t capacity;
    uint32_t usage;
    char *allocation;

    uint32_t maximum;
    ufn_arena_s *continuation;
};

/// @brief Allocates a new arena tree with up to 'max_tree_size' .
/// @warning No later but possibly earlier than when a total of
///         'max_tree_size'  was distributed, the allocation function
///         (ufn_arena_alloc) will return NULL.
/// @param capacity Number of bytes to allocate in the first arena;
///                 Determines how many bytes can be distributed.
/// @param max_tree_size Maximum number of bytes to allocate for the
///                      continuations. If this is less than or equal to
///                      'capacity' * 2, the next continuation will be
///                      the last one.
/// @return Newly created dynamic arena.
ufn_arena_s * ufn_new_dynamic_arena(uint32_t capacity, uint32_t max_tree_size);

/// @brief Allocates a new arena with an immutable capacity.
/// @warning If this arena's allocation is full,
///          'ufn_arena_alloc' will return NULL!
/// @note Internally calls 'ufn_new_dynamic_arena' with a maximum tree size of 0.
/// @param capacity Number of bytes to allocate in the arena;
///                 Determines how many bytes can be distributed.
/// @return Newly created arena.
ufn_arena_s * ufn_new_static_arena(uint32_t capacity);

/// @brief Frees an arena and all its dynamic continuations.
/// @param arena Arena to free.
void ufn_free_arena(ufn_arena_s *arena);

/// @brief Allocates a block of size 'size' in the arena 'arena'
/// @warning Will return NULL if no memory block of length 'size'
///          is available in the arena or any of the continuations.
/// @param arena Arena to use for allocating 'size' bytes.
/// @param size Number of bytes to allocate.
/// @return Pointer to the newly allocated block.
void * ufn_arena_alloc(ufn_arena_s *arena, uint32_t size);

/// @brief Clones a string into the arena.
/// @param arena Arena to clone string into.
/// @param string String to clone.
/// @return Copied string. The size is immutable.
char * ufn_arena_clone_string(ufn_arena_s *arena, const char *string);

// END OF API, START OF IMPLEMENTATION.

#ifdef UFN_IMPLEMENTATION
#include <stdlib.h>
#include <string.h>

ufn_arena_s * ufn_new_dynamic_arena(uint32_t capacity, uint32_t max_tree_size)
{
    ufn_arena_s *arena = malloc(sizeof(ufn_arena_s));
    arena->capacity = capacity;
    arena->usage = 0;
    arena->allocation = malloc(capacity);
    arena->maximum = max_tree_size;
    arena->continuation = NULL;
    return arena;
}
ufn_arena_s * ufn_new_static_arena(uint32_t capacity)
{
    return ufn_new_dynamic_arena(capacity, 0);
}

void ufn_free_arena(ufn_arena_s *arena)
{
    if(arena->continuation != NULL)
    {
        ufn_free_arena(arena->continuation);
    }
    free(arena->allocation);
    free(arena);
}

void * ufn_arena_alloc(ufn_arena_s *arena, uint32_t size)
{
    if((arena->usage + size) >= arena->capacity)
    {
        if(arena->continuation == NULL)
        {
            uint32_t continuation_size = arena->capacity * 2;
            // Clamp the continuation's size to the maximum allowed one.
            if(continuation_size > arena->maximum)
            {
                continuation_size = arena->maximum;
            }
            uint32_t continued_maximum = continuation_size - arena->maximum;
            // If no continuation of this arena tree can fulfill
            // a request of this size,
            // -> Deny the request and return NULL.
            if(size >= arena->maximum)
            {
                return NULL;
            }
            // If no more arena can be allocated (The tree has its maximum size)
            // -> return NULL.
            if(continuation_size == 0)
            {
                return NULL;
            }
            arena->continuation = ufn_new_dynamic_arena(continuation_size, continued_maximum);
        }
        return ufn_arena_alloc(arena->continuation, size);
    }
    void *block = arena->allocation + arena->usage;
    arena->usage += size;
    return block;
}

char * ufn_arena_clone_string(ufn_arena_s *arena, const char *string)
{
    uint32_t len_string = strlen(string);
    char *string_copy = ufn_arena_alloc(arena, len_string + 1);
    if(string_copy == NULL)
    {
        return NULL;
    }
    memcpy(string_copy, string, len_string + 1);
    return string_copy;
}

#endif // UFN_IMPLEMENTATION

#endif // UFN_ARENA_H