rntviewer/src/mem.cpp

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

#define align_pow2(x, b) (((x) + (b) - 1) & ( ~((b) - 1)))

#ifdef ENABLE_ASAN
#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

thread_local Thread_Ctx *thread_local_ctx = nullptr;

internal
void tctx_init(Thread_Ctx &tctx) {
  memset(&tctx, 0, sizeof(Thread_Ctx));
  Arena **arena_ptr = tctx.arenas;
  for (u64 i = 0; i < countof(tctx.arenas); ++i, ++arena_ptr) {
    *arena_ptr = arena_alloc();
  }
  thread_local_ctx = &tctx;
}

internal
void tctx_release()
{
  for (u64 i = 0; i < countof(thread_local_ctx->arenas); ++i)
    arena_release(thread_local_ctx->arenas[i]);
}

internal
Arena *tctx_get_scratch(Arena **conflicts, u64 count)
{
  Thread_Ctx *tctx = thread_local_ctx;
  Arena *res = nullptr;
  Arena **arena_ptr = tctx->arenas;
  for (u64 i = 0; i < countof(tctx->arenas); ++i, ++arena_ptr) {
    Arena **conflict_ptr = conflicts;
    b8 has_conflict = false;
    for (u64 j = 0; j < count; ++j, ++conflict_ptr) {
      if (*arena_ptr == *conflict_ptr) {
        has_conflict = true;
        break;
      }
    }
    if (!has_conflict) {
      res = *arena_ptr;
      break;
    }
  }

  return res;
}

// 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 = os_page_size();
  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;
    arena->mem_used = arena->mem_peak_used = ARENA_HEADER_SIZE;
  }

  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;
      new_block->prev = arena->cur;
      arena->cur = new_block;

      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;
    u64 added_size = pos_new - cur->pos;
    arena->mem_used += added_size;
    arena->mem_peak_used = max(arena->mem_peak_used, arena->mem_used);
    cur->pos = pos_new;
    asan_unpoison_memory_region(mem, added_size);
  }

  return mem;
}

template <typename T>
internal
T *arena_push_array_nozero(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_nozero<T>(arena, count);
  memset(ary, 0, sizeof(T) * count);
  return ary;
}

template <typename T>
internal
T *arena_push(Arena *arena)
{
  return arena_push_array<T>(arena, 1);
}

template <typename T>
internal
T *arena_push_nozero(Arena *arena)
{
  return arena_push_array_nozero<T>(arena, 1);
}

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 = max(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));

  arena->mem_used -= (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);
}