memory.h

#ifndef MEMORY_H
#define MEMORY_H

#include "language_layer.h"
#include <string.h>

struct Memory_Arena {
    u8 *base;
    usize used;
    usize total;
};

#define push_struct(arena, type) (type*)_push_size(arena, sizeof(type))
#define push_array(arena, count, type) (type*)_push_size(arena, sizeof(type) * count)
#define push_size(arena, size) (u8*)_push_size(arena, size)

inline void* _push_size(Memory_Arena* arena, usize size) {
    assert(arena && arena->used + size < arena->total);

    u8* result = arena->base + arena->used;
    arena->used += size;
    return result;
}

struct Temp_Memory {
    Memory_Arena* arena;
    usize used;
};

inline Temp_Memory begin_temp_memory(Memory_Arena* arena) {
    Temp_Memory result;

    result.arena = arena;
    result.used = arena->used;
    return result;
}

inline void end_temp_memory(Temp_Memory temp_mem) {
    Memory_Arena* arena = temp_mem.arena;
    assert(arena);
    arena->used = temp_mem.used;
}

struct Pool_Bucket {
    void* allocation;
    usize used;
};

struct Pool_Allocator {
    Pool_Bucket* buckets;
    u32 num_buckets;
    u32 bucket_size;
    u8* memory;
};

inline Pool_Allocator init_pool_allocator(Memory_Arena* arena, u32 num_buckets, u32 bucket_size) {
    Pool_Allocator result;
    result.num_buckets = num_buckets;
    result.bucket_size = bucket_size;

    assert(bucket_size % 2 == 0);

    result.buckets = push_array(arena, num_buckets, Pool_Bucket);
    result.memory = push_size(arena, num_buckets * bucket_size);

    memset(result.buckets, 0, num_buckets * sizeof(*result.buckets));

    return result;
}

#define pool_struct(pool, type) (type*)pool_alloc(pool, sizeof(type))
#define pool_array(pool, count, type) (type*)pool_alloc(pool, count * sizeof(type))
#define pool_size(pool, size) (u8*)pool_alloc(pool, size)

inline void* pool_alloc(Pool_Allocator* pool, usize size) {
    // @NOTE(CHall): Find enough sequential buckets for this allocation
    s32 start_index = -1;
    u32 space_found = 0;
    For (Pool_Bucket, pool->buckets, pool->num_buckets) {
        if (!it->allocation && start_index == -1) {
            start_index = (s32)(pool->buckets + pool->num_buckets - it);
        } else if (it->allocation) {
            start_index = -1;
            space_found = 0;
        }

        if (start_index != -1) ++space_found;

        if ((usize)(space_found * pool->bucket_size) >= size) break;
    }

    // @NOTE(CHall): If we couldn't find enough room then return nullptr
    if (start_index == -1 || (usize)(space_found * pool->bucket_size) < size) return nullptr;

    // @NOTE(CHall): Find the allocation in the raw memory
    u8* const allocation = pool->memory + (start_index * pool->bucket_size);

    // @NOTE(CHall): Update bucket data on pool allocator
    for (u32 i = 0; i < space_found; ++i) {
        Pool_Bucket* const bucket = pool->buckets + i + start_index;
        bucket->allocation = allocation;
        bucket->used = size > pool->bucket_size ? pool->bucket_size : size;
        size -= pool->bucket_size;
    }

    return allocation;
}

inline void pool_free(Pool_Allocator* pool, void* ptr) {
    // @NOTE(CHall): Find all the buckets that have ptr as their allocation and reset the data
    b32 found_allocation = false;
    For(Pool_Bucket, pool->buckets, pool->num_buckets) {
        if (it->allocation == ptr) found_allocation = true;

        if (found_allocation && it->allocation != ptr) break;

        // @NOTE(CHall): Reset bucket data
        if (found_allocation) {
            it->allocation = nullptr;
            it->used = 0;
        }
    }

    // @NOTE(CHall): Assert if we didn't find the allocation
    assert(found_allocation);
}

inline void* pool_realloc(Pool_Allocator* pool, void* ptr, usize size) {
    // @NOTE(CHall): Try to expand the allocation if we have the space. 
    // @TODO(CHall): We should set this up where we don't set up the buckets until we know we have the space.
    usize old_size = 0;
    b32 found_allocation = false;
    For (Pool_Bucket, pool->buckets, pool->num_buckets) {
        if (it->allocation == ptr) found_allocation = true;

        if (found_allocation && it->allocation && it->allocation != ptr) break;

        // @NOTE(CHall): Set expanded buckets with allocation data.         
        it->allocation = ptr;
        it->used = pool->bucket_size;
        if (old_size + pool->bucket_size >= size) {
            it->used = size - old_size;
            return it->allocation;
        }
        old_size += it->used;
    }

    assert(found_allocation && old_size < size);

    // @NOTE(CHall): If we couldn't find enough space just alloc and move memory
    void* result = pool_alloc(pool, size);
    memcpy(result, ptr, old_size);
    pool_free(pool, ptr);
    return result;
}

#endif /* MEMORY_H */