/*
    Memory Arena: A big pool of memory allocated once to allocate your own arrays from.
    Can guarantee contiguous memory for all data.

    Usage:

        MemoryArena *arena = new MemoryArena(size);
        T *T_ptr = arena->alloc_1D<T>(len, init);
        T **T_pptr = arena->alloc_2D<T>(len1, len2, init);

    Free all the memory using:

        delete arena;

    NOTE: WILL leave dangling pointers after destruction.
*/

struct MemoryArena {
    size_t arena_size;
    size_t current;
    char *memory;

    // use a triple pointer to store a reference to each double pointer, so that
    // the inner pointers can be deallocated
    void ***memory_double_pointers;
    size_t num_double_pointers;
    size_t double_pointer_counter;

    MemoryArena(size_t size, int num_double_pointers) : arena_size(size), num_double_pointers(num_double_pointers) {
        memory = new char[arena_size];
        for (size_t i = 0; i < arena_size; i++)
            memory[i] = 0;
        current = 0;
        printf("Allocate arena of size %d at 0x%p, memory byte alignment = %d\n", (int)size, memory, (int)((uintptr_t)memory & 15));

        double_pointer_counter = 0;
        memory_double_pointers = new void**[num_double_pointers];
        printf("Preparing to keep track of %d double pointers\n", (int)num_double_pointers);

    }
    ~MemoryArena() {
        printf("Destroying memory arena: %d double pointers and memory of size %d bytes\n", (int)double_pointer_counter, (int)arena_size);

        // destroying pointers from double pointers
        for (size_t i = 0; i < num_double_pointers; i++) {
            if (memory_double_pointers[i])
                delete[] memory_double_pointers[i];
        }
        // destryong double pointers
        delete[] memory_double_pointers;

        // destroying contigious memory
        delete[] memory;
    }

    void* alloc_block(size_t block_size) {
        size_t available_size = arena_size - current;

        if (available_size < block_size) {
            printf("Not enough space to allocate block of size %d, only %d left\n", (int)block_size, (int)available_size);
            return NULL;
        } else {
            size_t old_current = current;
            current = current + block_size;
            printf("Allocate %d bytes from arena, at index %d. Moving index to %d. %d left\n", (int)block_size, (int)old_current, (int)current, (int)(available_size-block_size));
            return memory + old_current;
        }
    }

    // typename T need to have operator= overloaded
    template <typename T>
    T* alloc_1D(size_t len, T init) {
        T *data = (T*)alloc_block(len*sizeof(T));
        if (data) {
            T *arr = data;
            for (size_t i = 0; i < len; i++)
                arr[i] = init;
            return arr;
        } else {
            return NULL;
        }
    }

    // For 2D arrays/double pointers we need to keep track of the
    // double pointers themselves so that we can free the memory used
    // for the inner pointers
    // these are stored in void ***memory_double_pointers
    template <typename T>
    T** alloc_2D(size_t len1, size_t len2, T init) {
        // allocate 1D block of contiguous memory
        T *data = (T*)alloc_block(len1*len2*sizeof(T));

        if (data) {
            T** arr = new T*[len1];
            memory_double_pointers[double_pointer_counter++] = (void**)arr;

            for (size_t i = 0; i < len1; i++) {
                arr[i] = data + len2*i;
                for (size_t j = 0; j < len2; j++) {
                    arr[i][j] = init;
                }
            }
            return arr;
        } else {
            return NULL;
        }
    }
};