#ifndef ARENA_H
#define ARENA_H

#include "memory.h"

#define arena_push(a, type) ((type *)arena_push_bytes((a), sizeof(type), alignof(type)))
#define arena_push_zero(a, type) ((type *)arena_push_bytes_zero((a), sizeof(type), alignof(type)))

#define arena_push_array(a, type, n) ((type *)arena_push_bytes((a), (sizeof(type) * (n)), alignof(type)))
#define arena_push_array_zero(a, type, n) ((type *)arena_push_bytes_zero((a), (sizeof(type) * (n)), alignof(type)))

#define arena_pop(a, type, dest) arena_pop_struct((a), sizeof(type), dest)
#define arena_pop_array(a, type, n, dest) arena_pop_struct((a), sizeof(type) * (n), dest)

/* Returns a pointer to where the next allocation would be (at alignment of type).
 * Equivalent arena_push but without actually allocating anything. */
#define arena_dry_push(a, type) (type *)(_arena_dry_push((a), alignof(type)))

struct temp_arena {
    struct arena *arena;
    u64 start_pos;

#if RTC
    u64 scratch_id;
#endif
};

struct arena arena_alloc(u64 reserve);
void arena_release(struct arena *arena);
void *arena_push_bytes(struct arena *arena, u64 size, u64 align);
void arena_copy_replace(struct arena *dest, struct arena *src);
void arena_decommit_unused_blocks(struct arena *arena);
void arena_set_readonly(struct arena *arena);
void arena_set_readwrite(struct arena *arena);

INLINE void *arena_push_bytes_zero(struct arena *arena, u64 size, u64 align)
{
    void *p = arena_push_bytes(arena, size, align);
    MEMZERO(p, size);
    return p;
}

INLINE void arena_pop_to(struct arena *arena, u64 pos)
{
    ASSERT(arena->pos >= pos);
    ASSERT(!arena->readonly);

    ASAN_POISON(arena->base + pos, arena->pos - pos);
    arena->pos = pos;
}

INLINE void arena_pop_struct(struct arena *arena, u64 size, void *copy_dest)
{
    ASSERT(arena->pos >= size);
    ASSERT(!arena->readonly);

    u64 new_pos = arena->pos - size;
    void *src = (void *)(arena->base + new_pos);
    MEMCPY(copy_dest, src, size);

    ASAN_POISON(arena->base + new_pos, arena->pos - new_pos);
    arena->pos = new_pos;
}

INLINE void *arena_align(struct arena *arena, u64 align)
{
    ASSERT(!arena->readonly);

    if (align > 0) {
        u64 aligned_start_pos = (arena->pos + (align - 1));
        aligned_start_pos -= aligned_start_pos % align;
        u64 align_bytes = aligned_start_pos - (u64)arena->pos;
        if (align_bytes > 0) {
            return (void *)arena_push_array(arena, u8, align_bytes);
        } else {
            return (void *)(arena->base + arena->pos);
        }
    } else {
        /* 0 alignment */
        ASSERT(false);
        return (void *)(arena->base + arena->pos);
    }
}

INLINE struct temp_arena arena_temp_begin(struct arena *arena)
{
    struct temp_arena t = ZI;
    t.arena = arena;
    t.start_pos = arena->pos;
    return t;
}

INLINE void arena_temp_end(struct temp_arena temp)
{
    arena_pop_to(temp.arena, temp.start_pos);
}

INLINE void arena_reset(struct arena *arena)
{
    arena_pop_to(arena, 0);
}

INLINE struct string arena_to_string(struct arena *arena)
{
    struct string b;
    b.text = arena->base;
    b.len = arena->pos;
    return b;
}

INLINE void *_arena_dry_push(struct arena *arena, u64 align)
{
    u64 aligned_start_pos = (arena->pos + (align - 1));
    aligned_start_pos -= aligned_start_pos % align;
    void *ptr = arena->base + aligned_start_pos;
    return ptr;
}

#endif