#include "malunal/allocators/stack.h"

typedef struct {
  allocator_t      allocator;
  allocator_mptr_t upstream;

  malunal_size_t stride;
  malunal_size_t count;
  malunal_size_t capacity;
  malunal_mptr_t buffer;
} impl_t;

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

static
error_t
stack_allocator_acquire_impl(
  stack_allocator_mptr_t allocator,
  malunal_size_t         size,
  malunal_mptr_t*        out
) {
  MALUNAL_UNUSED(size);
  return stack_allocator_acquire(allocator, out);
}

static
error_t
stack_allocator_dispose_impl(
  stack_allocator_mptr_t allocator,
  malunal_mptr_t         address,
  malunal_size_t         size
) {
  MALUNAL_UNUSED(size);
  return stack_allocator_dispose(allocator, address);
}

static const
allocator_vtable_t stack_allocator_vtable = {
  .acquire = (allocator_acquire_pfn_t)&stack_allocator_acquire_impl,
  .dispose = (allocator_dispose_pfn_t)&stack_allocator_dispose_impl
};

error_t
stack_allocator_stride(
  stack_allocator_iptr_t allocator,
  malunal_size_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;
  *out = self->stride;
  return NO_ERROR;
}

error_t
stack_allocator_count(
  stack_allocator_iptr_t allocator,
  malunal_size_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;
  *out = self->count;
  return NO_ERROR;
}

error_t
stack_allocator_capacity(
  stack_allocator_iptr_t allocator,
  malunal_size_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;
  *out = self->capacity;
  return NO_ERROR;
}

error_t
stack_allocator_init(
  malunal_size_t         stride,
  malunal_size_t         capacity,
  allocator_mptr_t       upstream,
  stack_allocator_mptr_t allocator
) {
  if (stride < 8)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_FAILURE
    };

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

  if (upstream == null)
    upstream = libc_allocator();

  malunal_mptr_t address;
  malunal_size_t bytes  = stride * capacity;
  error_t        result = allocator_acquire(upstream, bytes, &address);
  if (result.domain != null)
    return result;

  impl_mptr_t self = (impl_mptr_t)allocator;
  self->allocator = (allocator_t){ &stack_allocator_vtable };
  self->upstream  = upstream;
  self->stride    = stride;
  self->count     = 0;
  self->capacity  = capacity;
  self->buffer    = address;
  return NO_ERROR;
}

error_t
stack_allocator_acquire(
  stack_allocator_mptr_t allocator,
  malunal_mptr_t*        out
) {
  if (allocator == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NULL_ALLOCATOR
    };

  impl_mptr_t self = (impl_mptr_t)allocator;
  if (self->count == self->capacity)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_OUT_OF_MEMORY
    };

  malunal_uint8_t* asbytes = self->buffer;
  *out = asbytes + self->stride * self->count;
  self->count++;
  return NO_ERROR;
}

error_t
stack_allocator_dispose(
  stack_allocator_mptr_t allocator,
  malunal_mptr_t         address
) {
  if (allocator == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NULL_ALLOCATOR
    };

  impl_mptr_t      self = (impl_mptr_t)allocator;
  malunal_uint8_t* addr = address;
  malunal_uint8_t* beg  = self->buffer;
  malunal_uint8_t* end  = beg + self->stride * self->capacity;
  if (addr < beg || addr > end)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NOT_MY_ADDRESS
    };

  malunal_uint8_t* top = self->buffer;
  top += self->stride * (self->count - 1);
  if (addr != top)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_FAILURE
    };

  self->count--;
  return NO_ERROR;
}

error_t
stack_allocator_free(
  stack_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 = allocator_dispose(
    self->upstream,
    self->buffer,
    self->stride * self->capacity
  );
  if (result.domain != null)
    return result;

  self->stride   = 0;
  self->count    = 0;
  self->capacity = 0;
  self->buffer   = null;
  return NO_ERROR;
}