// Mostly taken from the raddebugger -- thanks, Ryan.

#define align_pow2(x, b) (((x) + (b) - 1) & ( ~((b) - 1)))
#define SLLStackPush_N(f,n,next) ((n)->next=(f), (f)=(n))
#define SLLStackPop_N(f,next) ((f)=(f)->next)

#ifdef DEBUG
#define asan_poison_memory_region(mem, cmt) __asan_poison_memory_region(mem, cmt)
#define asan_unpoison_memory_region(mem, cmt) __asan_unpoison_memory_region(mem, cmt)
#else
#define asan_poison_memory_region(mem, cmt)
#define asan_unpoison_memory_region(mem, cmt)
#endif

// @Platform
internal 
void *os_reserve(u64 size)
{
  return mmap(0, size, PROT_NONE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
}

// @Platform
internal 
void os_release(void *mem, u64 size)
{
  munmap(mem, size);
}

// @Platform
internal 
b32x os_commit(void *addr, u64 size)
{
  return mprotect(addr, size, PROT_READ|PROT_WRITE) == 0;
}

// init_res: initial reserved size
// init_cmt: initial committed size
internal 
Arena *arena_alloc_sized(u64 init_res, u64 init_cmt)
{
  assert(ARENA_HEADER_SIZE < init_cmt && init_cmt <= init_res);

  u64 page_size = getpagesize();
  u64 res = align_pow2(init_res, page_size);
  u64 cmt = align_pow2(init_cmt, page_size);

  // reserve memory
  void *mem = os_reserve(res);
  if (!os_commit(mem, cmt)) {
    munmap(mem, init_res);
  }

  Arena *arena = (Arena *)mem;
  if (arena) {
    asan_poison_memory_region(mem, cmt);
    asan_unpoison_memory_region(mem, ARENA_HEADER_SIZE);

    arena->cur = arena;
    arena->pos = ARENA_HEADER_SIZE;
    arena->cmt = cmt;
    arena->res = res;
    arena->align = 8;
    arena->grow = true;
  }

  return arena;
}

internal 
Arena *arena_alloc()
{
  return arena_alloc_sized(ARENA_RESERVE_SIZE, ARENA_COMMIT_SIZE);
}

internal 
void arena_release(Arena *arena)
{
  for (Arena *node = arena->cur, *prev = 0; node; node = prev) {
    prev = node->prev;
    os_release(node, node->res);
  }
}

internal
u64 arena_huge_push_threshold()
{
  u64 res = ARENA_RESERVE_SIZE;
  u64 threshold = (res - ARENA_HEADER_SIZE) / 2 + 1;
  return threshold;
}

internal
void *arena_push_impl(Arena *arena, u64 size)
{
  Arena *cur = arena->cur;
  u64 pos_mem = align_pow2(cur->pos, arena->align);
  u64 pos_new = pos_mem + size;

  if (cur->res < pos_new && arena->grow) {
    Arena *new_block;
    if (size < arena_huge_push_threshold()) {
      new_block = arena_alloc();
    } else {
      u64 new_block_size = size + ARENA_HEADER_SIZE;
      new_block = arena_alloc_sized(new_block_size, new_block_size);
    }

    if (new_block) {
      new_block->base_pos = cur->base_pos + cur->res;
      SLLStackPush_N(arena->cur, new_block, prev);

      cur = new_block;
      pos_mem = align_pow2(cur->pos, cur->align);
      pos_new = pos_mem + size;
    }
  }

  if (cur->cmt < pos_new) {
    u64 cmt_new_aligned = align_pow2(pos_new, ARENA_COMMIT_SIZE);
    u64 cmt_new_clamped = min(cmt_new_aligned, cur->res);
    u64 cmt_new_size = cmt_new_clamped - cur->cmt;
    b32x is_cmt_ok = os_commit((u8*)cur + cur->cmt, cmt_new_size);

    if (is_cmt_ok)
      cur->cmt = cmt_new_clamped;
  }

  void *mem = 0;
  if (cur->cmt >= pos_new) {
    mem = (u8*)cur + pos_mem;
    cur->pos = pos_new;
    asan_unpoison_memory_region(mem, size);
  }

  return mem;
}

template <typename T>
internal
T *arena_push(Arena *arena)
{
  void *mem = arena_push_impl(arena, sizeof(T));
  return (T *)mem;
}

template <typename T>
internal
T *arena_push_array_no_zero(Arena *arena, u64 count)
{
  void *mem = arena_push_impl(arena, sizeof(T) * count);
  return (T *)mem;
}

template <typename T>
internal
T *arena_push_array(Arena *arena, u64 count)
{
  T *ary = arena_push_array_no_zero<T>(arena, count);
  memset(ary, 0, sizeof(T) * count);
  return ary;
}

internal
u8 *arena_push_contiguous(Arena *arena, u64 size)
{
  b32 restore = arena->grow;
  arena->grow = 0;
  void *mem = arena_push_impl(arena, size);
  arena->grow = restore;
  return (u8 *)mem;
}

internal
u64 arena_pos(Arena *arena)
{
  Arena *cur = arena->cur;
  u64 pos = cur->base_pos + cur->pos;
  return pos;
}

internal
void arena_pop_to(Arena *arena, u64 big_pos_unclamped)
{
  u64 big_pos = min(ARENA_HEADER_SIZE, big_pos_unclamped);
  
  // unroll the chain
  Arena *current = arena->cur;
  for (Arena *prev = 0; current->base_pos >= big_pos; current = prev) {
    prev = current->prev;
    os_release(current, current->res);
  }
  assert(current);
  arena->cur = current;
  
  // compute arena-relative position
  u64 new_pos = big_pos - current->base_pos;
  assert(new_pos <= current->pos);
  
  // poison popped memory block
  asan_poison_memory_region((u8*)current + new_pos, (current->pos - new_pos));
  
  // update position
  current->pos = new_pos;
}

internal
Temp temp_begin(Arena *arena)
{
  u64 pos = arena_pos(arena);
  Temp temp = {arena, pos};
  return temp;
}

internal
void temp_end(Temp temp)
{
  arena_pop_to(temp.arena, temp.pos);
}