#include <memory.h>
#include "malunal/allocators/arena.h"

typedef struct {
  malunal_uint32_t size;
  malunal_uint32_t used;
  malunal_mptr_t   next;
  malunal_uint8_t  bytes[];
} region_t;

typedef       region_t* region_mptr_t;
typedef const region_t* region_iptr_t;

typedef struct {
  allocator_t   allocator;
  region_mptr_t context;
} impl_t;

typedef       impl_t* impl_mptr_t;
typedef const impl_t* impl_iptr_t;
_Static_assert(
  sizeof(arena_allocator_t) == sizeof(impl_t),
  "Arena allocator must be the size of its implementation"
);

// Every allocation is rounded to this alignment so that any object may be stored
// within it. The region header is also this size, so offsets stay aligned.
#define ARENA_ALIGNMENT 16

// Region sizes are tracked in 32 bits, so nothing larger may be allocated.
#define ARENA_MAX_REQUEST ((malunal_uint32_t)-1 - sizeof(region_t))

static
error_t
arena_out_of_memory() {
  return (error_t) {
    .domain = &ERROR_DOMAIN_ALLOCATOR_T,
    .code   = ALLOCATOR_ERROR_OUT_OF_MEMORY
  };
}

static
error_t
create_region(
  malunal_size_t minimum,
  region_mptr_t* region
) {
  if (minimum > ARENA_MAX_REQUEST)
    return arena_out_of_memory();

  // Requests larger than the default region size get a region of their own.
  malunal_size_t wanted = sizeof(region_t) + minimum;
  if (wanted < MALUNAL_ARENA_REGION_SIZE)
    wanted = MALUNAL_ARENA_REGION_SIZE;

  malunal_size_t rounded = align_to_page(wanted);
  if (rounded < wanted || rounded > (malunal_uint32_t)-1)
    return arena_out_of_memory();

  allocator_mptr_t allocator = platform_allocator();
  malunal_mptr_t*  address   = (malunal_mptr_t*)region;
  error_t          result    = allocator_acquire(allocator, rounded, address);

  if (result.domain != null)
    return result;

  (*region)->size = rounded;
  (*region)->used = sizeof(region_t);
  (*region)->next = null;
  return NO_ERROR;
}

static
error_t
delete_region(region_mptr_t region) {
  if (region == null)
    return NO_ERROR;

  error_t result = delete_region(region->next);
  if (result.domain != null)
    return result;

  region->next = null;
  allocator_mptr_t allocator = platform_allocator();
  return allocator_dispose(allocator, region, region->size);
}

static
error_t
region_acquire(
  region_mptr_t   region,
  malunal_size_t  size,
  malunal_mptr_t* out
) {
  malunal_uint32_t offset =
    region->used - sizeof(region_t);
  *out = region->bytes + offset;
  region->used += size;
  return NO_ERROR;
}

static const
allocator_vtable_t arena_allocator_vtable = {
  .acquire = (allocator_acquire_pfn_t)&arena_allocator_acquire,
  .dispose = (allocator_dispose_pfn_t)&arena_allocator_dispose
};

error_t
arena_allocator_size(
  arena_allocator_iptr_t allocator,
  malunal_uint32_t*      out
) {
  if (allocator == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NULL_ALLOCATOR
    };

  impl_iptr_t    self   = (impl_iptr_t)allocator;
  region_iptr_t  region = self->context;
  malunal_size_t total  = 0;
  while (region != null) {
    total += region->size;
    region = region->next;
  }

  *out = total;
  return NO_ERROR;
}

error_t
arena_allocator_used(
  arena_allocator_iptr_t allocator,
  malunal_uint32_t*      out
) {
  if (allocator == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NULL_ALLOCATOR
    };

  impl_iptr_t    self   = (impl_iptr_t)allocator;
  region_iptr_t  region = self->context;
  malunal_size_t total  = 0;
  while (region != null) {
    total += region->used;
    region = region->next;
  }

  *out = total;
  return NO_ERROR;
}

error_t
arena_allocator_init(
  malunal_size_t         capacity,
  arena_allocator_mptr_t allocator
) {
  if (allocator == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NULL_ALLOCATOR
    };

  region_mptr_t region = null;
  error_t       result = create_region(capacity, &region);
  if (result.domain != null)
    return result;

  impl_mptr_t self = (impl_mptr_t)allocator;
  self->allocator = (allocator_t){ &arena_allocator_vtable };
  self->context   = region;
  return NO_ERROR;
}

error_t
arena_allocator_acquire(
  arena_allocator_mptr_t allocator,
  malunal_size_t         size,
  malunal_mptr_t*        out
) {
  if (allocator == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NULL_ALLOCATOR
    };

  if (size > ARENA_MAX_REQUEST)
    return arena_out_of_memory();
  size = align_to(size, ARENA_ALIGNMENT);

  impl_mptr_t   self   = (impl_mptr_t)allocator;
  region_mptr_t region = self->context;
  region_mptr_t last   = null;
  while (region != null) {
    if (region->size - region->used >= size)
      return region_acquire(region, size, out);
    last   = region;
    region = region->next;
  }

  region_mptr_t created = null;
  error_t       result  = create_region(size, &created);
  if (result.domain != null)
    return result;

  if (last == null)
    self->context = created;
  else
    last->next = created;
  return region_acquire(created, size, out);
}

error_t
arena_allocator_dispose(
  arena_allocator_mptr_t allocator,
  malunal_mptr_t         address,
  malunal_size_t         size
) {
  if (allocator == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NULL_ALLOCATOR
    };

  // Arenas never release individual allocations, disposing only verifies that
  // the address belongs to one of the regions.
  impl_iptr_t      self = (impl_iptr_t)allocator;
  malunal_uint8_t* addr = address;
  for (region_iptr_t reg = self->context; reg != null; reg = reg->next) {
    const malunal_uint8_t* beg = reg->bytes;
    const malunal_uint8_t* end = reg->bytes + reg->size - sizeof(region_t);
    if (addr >= beg && addr <= end)
      return NO_ERROR;
  }

  return (error_t) {
    .domain = &ERROR_DOMAIN_ALLOCATOR_T,
    .code   = ALLOCATOR_ERROR_NOT_MY_ADDRESS
  };
}

error_t
arena_allocator_reset(
  arena_allocator_mptr_t allocator
) {
  if (allocator == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NULL_ALLOCATOR
    };

  impl_mptr_t   self   = (impl_mptr_t)allocator;
  region_mptr_t region = self->context;
  while (region != null) {
    region->used = sizeof(region_t);
    region = region->next;
  }

  return NO_ERROR;
}

error_t
arena_allocator_free(
  arena_allocator_mptr_t allocator
) {
  if (allocator == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NULL_ALLOCATOR
    };

  impl_mptr_t self   = (impl_mptr_t)allocator;
  error_t     result = delete_region(self->context);
  if (result.domain != null)
    return result;

  self->context = null;
  return NO_ERROR;
}