#include "arena.h"
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <stdio.h>

#ifdef ARENA_ASAN_INSTRUMENT
#include <sanitizer/asan_interface.h>

#define ASAN_POISON(...) __asan_poison_memory_region(__VA_ARGS__)
#define ASAN_UNPOISON(...) __asan_unpoison_memory_region(__VA_ARGS__)
#else
#define ASAN_POISON(...)
#define ASAN_UNPOISON(...)
#endif

struct vl__arena_tag {
  size_t cap; /* number of bytes after the end of this struct we can use */
  size_t cur;
};

#define ARENA_ALIGN (size_t)((1 << 6) - 1)
#define ROUND_UP(_am) (((_am) + ARENA_ALIGN) & ~ARENA_ALIGN)

static void *arena_addr(void *a, size_t len)
{
  return (void *)((uintptr_t)a + ROUND_UP(sizeof(vl_arena)) + len);
}

/* returns NULL if an arena could not be allocated */
vl_arena *vl_arena_new(size_t cap)
{
  cap = ROUND_UP(cap);

  vl_arena *arena = calloc(1, ROUND_UP(sizeof(vl_arena)) + cap);
  if (!arena) return NULL;

  arena->cap = cap;
  arena->cur = 0;

  ASAN_POISON(arena + 1, cap + (ROUND_UP(sizeof(vl_arena)) - sizeof(vl_arena)));

  return arena;
}

/* aborts if the arena is overflowing */
void *vl_arena_push(vl_arena *parena, size_t len)
{
  size_t rlen = ROUND_UP(len);
  if (parena->cur + rlen > parena->cap) abort();

  void *ret = arena_addr(parena, parena->cur);
  parena->cur += rlen;

  ASAN_UNPOISON(ret, len);
  return ret;
}

/* resets the arena (but does not free it) */
void vl_arena_reset(vl_arena *parena)
{
  parena->cur = 0;
  ASAN_POISON(arena_addr(parena, 0), parena->cap);
}

/* frees the arena */
void vl_arena_free(vl_arena *parena)
{
  free(parena);
}

char *vl_arena_strdup(vl_arena *parena, const char *str)
{
  size_t len = strlen(str);
  void *out = vl_arena_push(parena, len + 1);
  memcpy(out, str, len + 1);
  return out;
}

char *vl_arena_sprintf(vl_arena *parena, const char *fmt, ...)
{
  va_list ap;
  va_start(ap, fmt);

  char *ret = vl_arena_vsprintf(parena, fmt, ap);

  va_end(ap);
  return ret;
}

char *vl_arena_vsprintf(vl_arena *parena, const char *fmt, va_list ap)
{
  va_list copy;
  va_copy(copy, ap);

  int desire = vsnprintf(NULL, 0, fmt, copy);
  va_end(copy);

  if (desire < 0) {
    abort(); /* vsnprintf should never fail */
  }

  char *str = vl_arena_push(parena, (unsigned)desire + 1);
  vsnprintf(str, (unsigned)desire + 1, fmt, ap);

  return str;
}