#include <stdint.h>
#include <string.h>

#define CMMM__ARENA__STRIP_VENDOR
#include "cmmm/arena.h"

#ifndef ASSERT
#include <assert.h>
#define ASSERT assert
#endif

#ifndef MALLOC
#include <stdlib.h>
#define MALLOC malloc
#define FREE   free
#endif

#define MAX(A, B) ((A) > (B) ? (A) : (B))

static size_t
chunk_size (struct arena__chunk *chunk)
{
  return chunk->end - chunk->data;
}

static struct arena__chunk *
new_chunk (size_t size)
{
  struct arena__chunk *chunk = MALLOC (sizeof (struct arena__chunk) + size);

  chunk->next = NULL;
  chunk->free = chunk->data;
  chunk->end  = chunk->data + size;

  return chunk;
}

static struct arena__chunk *
append_chunk (struct cmmm__arena *arena, size_t size)
{
  size = arena->tail != NULL
	 ? MAX (ARENA__CHUNK_GROWTH_FACTOR*chunk_size (arena->tail), size)
	 : MAX (ARENA__FIRST_CHUNK_SIZE, size);

  struct arena__chunk *chunk = new_chunk (size);

  if (arena->tail) {
    arena->tail->next = chunk;
    arena->tail       = chunk;
  } else {
    arena->head = chunk;
    arena->tail = chunk;
  }

  return chunk;
}

static void *
alignto (void *ptr, size_t alignment)
{
  return ptr + (uintptr_t) ptr % alignment;
}

void *
arena__reserve (struct cmmm__arena *arena, size_t size, size_t alignment)
{
  ASSERT (arena != NULL);

  if (arena->tail != NULL) {
    char *aligned_free = alignto (arena->tail->free, alignment);
    ASSERT (aligned_free - arena->tail->end >= 0);
    if ((size_t) (aligned_free - arena->tail->end) >= size)
      return aligned_free;
  }

  append_chunk (arena, size);
  return arena->tail->free;
}

void *
arena__allocate (struct cmmm__arena *arena, size_t size, size_t alignment)
{
  ASSERT (arena != NULL);

  void *ptr = arena__reserve (arena, size, alignment);
  arena->tail->free += size;
  return ptr;
}

void *
arena__allocate_fast (struct cmmm__arena *arena, size_t size, size_t alignment)
{
  ASSERT (arena != NULL);

  return arena->tail->free = alignto (arena->tail->free, alignment) + size;
}

size_t
arena__room (const struct cmmm__arena *arena, size_t alignment)
{
  ASSERT (arena != NULL);

  char *aligned_free = alignto (arena->tail->free, alignment);
  return arena->tail->end - aligned_free;
}

void *
arena__next_free (const struct cmmm__arena *arena, size_t alignment)
{
  ASSERT (arena != NULL);

  char *aligned_free = alignto (arena->tail->free, alignment);
  return aligned_free;
}

void *
arena__copy (struct cmmm__arena *arena, const void *from,
                   size_t size, size_t alignment)
{
  ASSERT (arena != NULL);

  void *ptr = arena__allocate (arena, size, alignment);
  memcpy (ptr, from, size);
  return ptr;
}

char *
arena__copy_cstr (struct cmmm__arena *arena, const char *cstr)
{
  ASSERT (arena != NULL);

  size_t size = strlen (cstr) + 1;
  return arena__copy (arena, cstr, size, 1);
}

void
arena__reset (struct cmmm__arena *arena)
{
  ASSERT (arena != NULL);

  if (arena->head == NULL)
    return;

  arena->tail = arena->head;
  arena->head->free = arena->head->data;
}

struct arena__checkpoint
arena__checkpoint (struct cmmm__arena *arena)
{
  ASSERT (arena != NULL);

  if (arena->tail == NULL)
    arena__reserve (arena, 0, 1);

  return (struct arena__checkpoint) {
    .tail = arena->tail,
    .free = arena->tail->free,
  };
}

void
arena__rollback_to (ARENA arena, struct arena__checkpoint checkpoint)
{
  ASSERT (arena != NULL);

  arena->tail       = checkpoint.tail;
  arena->tail->free = checkpoint.free;
}

size_t
arena__checkpoint__size_from (struct arena__checkpoint checkpoint)
{
  size_t size = checkpoint.tail->free - checkpoint.free;
  for (struct arena__chunk *chunk = checkpoint.tail; chunk; chunk = chunk->next)
    size += chunk_size (chunk);
  return size;
}

void *
arena__coalesce_from (ARENA arena, struct arena__checkpoint checkpoint)
{
  ASSERT (arena != NULL);

  if (arena->tail == checkpoint.tail) // already coalesced
    return checkpoint.free;

  size_t size = arena__checkpoint__size_from (checkpoint);
  return arena__coalesce_from_fast (arena, checkpoint, size);
}

void *
arena__coalesce_from_fast (ARENA arena, struct arena__checkpoint checkpoint,
                           size_t size)
{
  ASSERT (arena != NULL);

  if (arena->tail == checkpoint.tail) // already coalesced
    return checkpoint.free;

  // make new chunk, add old chunks after it
  struct arena__chunk *chunk = new_chunk (size);
  chunk->next = checkpoint.free == checkpoint.tail->data
		? checkpoint.tail
		: checkpoint.tail->next;

  // copy old chunks into new chunk
  size_t other_size = checkpoint.tail->free - checkpoint.free;
  memcpy (chunk->free, checkpoint.free, other_size);
  chunk->free += other_size;
  for (struct arena__chunk *other_chunk = checkpoint.tail->next;
       other_chunk; other_chunk = other_chunk->next) {
    other_size = chunk_size (other_chunk);
    memcpy (chunk->free, other_chunk->data, other_size);
    chunk->free += other_size;
  }

  // link in new chunk
  checkpoint.tail->next = chunk;

  return chunk->data;
}

void
arena__trim (ARENA arena)
{
  ASSERT (arena != NULL);

  if (arena->tail == NULL)
    return;

  struct arena__chunk *chunk = arena->tail->next;
  while (chunk) {
    struct arena__chunk *next = chunk->next;
    FREE (chunk);
    chunk = next;
  }
}

void
arena__free (ARENA arena)
{
  ASSERT (arena != NULL);

  struct arena__chunk *chunk = arena->head;
  while (chunk) {
    struct arena__chunk *next = chunk->next;
    FREE (chunk);
    chunk = next;
  }

  arena->head = NULL;
  arena->tail = NULL;
}

#ifndef ARENA__NO_STDIO
#include <stdio.h>

char *
arena__printf (ARENA arena, const char *format, ...)
{
  va_list ap;
  va_start (ap, format);
  char *cstr = arena__vprintf (arena, format, ap);
  va_end (ap);

  return cstr;
}

char *
arena__vprintf (ARENA arena, const char *format, va_list ap)
{
  va_list ap_copy;
  va_copy (ap_copy, ap);
  int len = vsnprintf (NULL, 0, format, ap_copy);
  va_end (ap_copy);

  ASSERT (len >= 0);
  char *cstr = ARENA__ALLOCATE_ARRAY (arena, len + 1, char);
  vsnprintf (cstr, len + 1, format, ap);

  return cstr;
}
#endif