#include "malunal/allocators/pool.h"

typedef struct free_t free_t;
typedef        free_t* free_mptr_t;
typedef const  free_t* free_iptr_t;

 struct free_t {
  free_mptr_t next;
};

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;
  free_mptr_t    freelist;
} impl_t;

typedef       impl_t* impl_mptr_t;
typedef const impl_t* impl_iptr_t;
_Static_assert(
  sizeof(pool_allocator_t) == sizeof(impl_t),
  "Pool allocator must be the size of its implemenation"
);

static
error_t
pool_allocator_acquire_impl(
  pool_allocator_mptr_t allocator,
  malunal_size_t        size,
  malunal_mptr_t*       out
) {
  MALUNAL_UNUSED(size);
  return pool_allocator_acquire(allocator, out);
}

static
error_t
pool_allocator_dispose_impl(
  pool_allocator_mptr_t allocator,
  malunal_mptr_t        address,
  malunal_size_t        size
) {
  MALUNAL_UNUSED(size);
  return pool_allocator_dispose(allocator, address);
}

static const
allocator_vtable_t pool_allocator_vtable = {
  .acquire = (allocator_acquire_pfn_t)&pool_allocator_acquire_impl,
  .dispose = (allocator_dispose_pfn_t)&pool_allocator_dispose_impl
};

error_t
pool_allocator_stride(
  pool_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
pool_allocator_count(
  pool_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
pool_allocator_capacity(
  pool_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
pool_allocator_init(
  malunal_size_t        stride,
  malunal_size_t        capacity,
  allocator_mptr_t      upstream,
  pool_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;

  malunal_size_t   index;
  malunal_uint8_t* asbytes = (malunal_uint8_t*)address;
  free_mptr_t      chunk   = (free_mptr_t)address;
  for (index = 1; index < capacity; index++)
    chunk = chunk->next = (free_mptr_t)(asbytes + stride * index);

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

error_t
pool_allocator_acquire(
  pool_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
    };

  free_mptr_t prev = self->freelist;
  free_mptr_t curr = prev;
  free_mptr_t next = prev->next;
  while (next != null) {
    prev = curr;
    curr = next;
    next = next->next;
  }

  *out = curr;
  self->count++;
  if (curr == self->freelist)
    self->freelist = null;
  else
    prev->next = null;
  return NO_ERROR;
}

error_t
pool_allocator_dispose(
  pool_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
    };

  if (self->freelist == null) {
    self->freelist       = (free_mptr_t)addr;
    self->freelist->next = null;
    self->count--;
    return NO_ERROR;
  }

  free_mptr_t chunk = self->freelist;
  while (chunk->next != null)
    chunk = chunk->next;
  chunk->next       = (free_mptr_t)addr;
  chunk->next->next = null;
  self->count--;
  return NO_ERROR;
}

error_t
pool_allocator_free(
  pool_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;
  self->freelist = null;
  return NO_ERROR;
}