#include <memory.h>
#include "malunal/containers/map.h"


#define MAP_CONTROL_EMPTY ((malunal_uint8_t)0x00)
#define MAP_CONTROL_ERASED ((malunal_uint8_t)0x01)
#define MAP_CONTROL_FILLED ((malunal_uint8_t)0x02)
#define MAP_LOAD_FACTOR 0.75f
#define MAP_NO_SLOT ((malunal_size_t)-1)
#define MAP_HASH_BASIS 0xCBF29CE484222325ULL
#define MAP_HASH_PRIME 0x00000100000001B3ULL


typedef struct {
  container_t      container;
  allocator_mptr_t allocator;
  malunal_size_t   keystride;
  malunal_size_t   valstride;
  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(map_container_t) == sizeof(impl_t),
  "Map container must be the size of its implementation"
);


static const
container_vtable_t map_container_vtable = {
  .allocator = (container_allocator_pfn_t)&map_container_allocator,
  .stride    = (container_stride_pfn_t)&map_container_stride,
  .count     = (container_count_pfn_t)&map_container_count,
  .capacity  = (container_capacity_pfn_t)&map_container_capacity,
  .resize    = (container_resize_pfn_t)&map_container_resize,
  .reserve   = (container_reserve_pfn_t)&map_container_reserve,
  .append    = (container_append_pfn_t)&map_container_append,
  .remove    = (container_remove_pfn_t)&map_container_remove,
  .contains  = (container_contains_pfn_t)&map_container_contains,
  .clear     = (container_clear_pfn_t)&map_container_clear
};

static
malunal_size_t
power_2_ceil(malunal_size_t capacity) {
  if (capacity <= 1)
    return 1;

  capacity--;
  malunal_size_t power = 2;
  while (capacity >>= 1)
    power <<= 1;
  return power;
}

static
malunal_size_t
map_container_bytes(
  malunal_size_t keystride,
  malunal_size_t valstride,
  malunal_size_t capacity
) {
  return capacity + (keystride + valstride) * capacity;
}

static
malunal_uint8_t*
map_container_control(impl_iptr_t self) {
  return (malunal_uint8_t*)self->buffer;
}

static
malunal_uint8_t*
map_container_key(
  impl_iptr_t    self,
  malunal_size_t index
) {
  malunal_uint8_t* slots =
    map_container_control(self) +
    self->capacity;
  return slots + (self->keystride + self->valstride) * index;
}

static
malunal_uint8_t*
map_container_value(
  impl_iptr_t    self,
  malunal_size_t index
) {
  return
    map_container_key(self, index) +
    self->keystride;
}

static
malunal_size_t
map_container_hash(
  malunal_iptr_t key,
  malunal_size_t keystride
) {
  malunal_uint8_t* bytes = key;
  malunal_uint64_t hash  = MAP_HASH_BASIS;
  for (malunal_size_t index = 0; index < keystride; index++) {
    hash ^= (malunal_uint64_t)bytes[index];
    hash *= MAP_HASH_PRIME;
  }

  return (malunal_size_t)hash;
}

static
malunal_bool_t
map_container_find(
  impl_iptr_t     self,
  malunal_iptr_t  key,
  malunal_size_t* outidx
) {
  *outidx = MAP_NO_SLOT;
  if (self->buffer == null || self->capacity == 0)
    return false;

  malunal_uint8_t* control = map_container_control(self);
  malunal_size_t   mask    = self->capacity - 1;
  malunal_size_t   index   = map_container_hash(key, self->keystride) & mask;
  malunal_size_t   vacancy = MAP_NO_SLOT;
  malunal_size_t   probed  = 0;
  while (probed < self->capacity) {
    if (control[index] == MAP_CONTROL_EMPTY) {
      *outidx = vacancy != MAP_NO_SLOT ? vacancy : index;
      return false;
    }

    if (control[index] == MAP_CONTROL_ERASED) {
      if (vacancy == MAP_NO_SLOT)
        vacancy = index;
    } else if (memcmp(map_container_key(self, index), key, self->keystride) == 0) {
      *outidx = index;
      return true;
    }

    index   = (index + 1) & mask;
    probed += 1;
  }

  *outidx = vacancy;
  return false;
}

static
malunal_void_t
map_container_place(
  impl_mptr_t    self,
  malunal_size_t index,
  malunal_iptr_t key,
  malunal_iptr_t value
) {
  map_container_control(self)[index] = MAP_CONTROL_FILLED;
  memcpy(map_container_key(self, index), key, self->keystride);
  if (self->valstride != 0)
    memcpy(map_container_value(self, index), value, self->valstride);
}

static
error_t
map_container_realloc(
  impl_mptr_t    self,
  malunal_size_t capacity
) {
  if (capacity <= self->capacity)
    return NO_ERROR;

  malunal_mptr_t buffer;
  malunal_size_t newcap = power_2_ceil(capacity);
  malunal_size_t bytes  = map_container_bytes(
    self->keystride,
    self->valstride,
    newcap
  );

  error_t result = allocator_acquire(self->allocator, bytes, &buffer);
  if (result.domain != null)
    return result;

  memset(buffer, MAP_CONTROL_EMPTY, newcap);

  impl_t rehashed = *self;
  rehashed.buffer   = buffer;
  rehashed.capacity = newcap;

  malunal_uint8_t* control = map_container_control(self);
  for (malunal_size_t index = 0; index < self->capacity; index++) {
    if (control[index] != MAP_CONTROL_FILLED)
      continue;

    malunal_size_t   slot;
    malunal_uint8_t* key = map_container_key(self, index);
    map_container_find(&rehashed, key, &slot);
    map_container_place(&rehashed, slot, key, map_container_value(self, index));
  }

  if (self->buffer != null) {
    bytes = map_container_bytes(
      self->keystride,
      self->valstride,
      self->capacity
    );

    result = allocator_dispose(self->allocator, self->buffer, bytes);
  }

  self->buffer   = buffer;
  self->capacity = newcap;
  return result;
}

error_t
map_container_init(
  malunal_size_t       keystride,
  malunal_size_t       valstride,
  allocator_mptr_t     allocator,
  map_container_mptr_t map
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  impl_mptr_t    self     = (impl_mptr_t)map;
  malunal_size_t capacity = power_2_ceil(MALUNAL_CONTAINERS_MAP_INIT_CAPACITY);
  malunal_size_t bytes    = map_container_bytes(keystride, valstride, capacity);
  error_t        result   = allocator_acquire(allocator, bytes, &self->buffer);
  if (result.domain != null)
    return result;

  memset(self->buffer, MAP_CONTROL_EMPTY, capacity);
  self->container = (container_t){ &map_container_vtable };
  self->allocator = allocator;
  self->keystride = keystride;
  self->valstride = valstride;
  self->count     = 0;
  self->capacity  = capacity;
  return NO_ERROR;
}

error_t
map_container_free(
  map_container_mptr_t map
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  impl_mptr_t self   = (impl_mptr_t)map;
  error_t     result = allocator_dispose(
    self->allocator,
    self->buffer,
    map_container_bytes(self->keystride, self->valstride, self->capacity)
  );

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

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

error_t
map_container_allocator(
  map_container_iptr_t map,
  allocator_mptr_t*    out
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  impl_iptr_t self = (impl_iptr_t)map;
  *out = self->allocator;
  return NO_ERROR;
}

error_t
map_container_stride(
  map_container_iptr_t map,
  malunal_size_t*      out
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  impl_iptr_t self = (impl_iptr_t)map;
  *out = self->keystride + self->valstride;
  return NO_ERROR;
}

error_t
map_container_key_stride(
  map_container_iptr_t map,
  malunal_size_t*      out
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  impl_iptr_t self = (impl_iptr_t)map;
  *out = self->keystride;
  return NO_ERROR;
}

error_t
map_container_value_stride(
  map_container_iptr_t map,
  malunal_size_t*      out
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  impl_iptr_t self = (impl_iptr_t)map;
  *out = self->valstride;
  return NO_ERROR;
}

error_t
map_container_count(
  map_container_iptr_t map,
  malunal_size_t*      out
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  impl_iptr_t self = (impl_iptr_t)map;
  *out = self->count;
  return NO_ERROR;
}

error_t
map_container_capacity(
  map_container_iptr_t map,
  malunal_size_t*      out
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  impl_iptr_t self = (impl_iptr_t)map;
  *out = self->capacity;
  return NO_ERROR;
}

error_t
map_container_keys(
  map_container_iptr_t map,
  container_mptr_t     container
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  if (container == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_RECEIVER
    };

  malunal_size_t stride = 0;
  impl_iptr_t    self   = (impl_iptr_t)map;
  error_t        result = container_stride(container, &stride);
  if (result.domain != null)
    return result;

  if (stride != self->keystride)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_FAILURE
    };

  malunal_uint8_t* control = map_container_control(self);
  for (malunal_size_t index = 0; index < self->capacity; index++) {
    if (control[index] != MAP_CONTROL_FILLED)
      continue;

    result = container_append(container, map_container_key(self, index));
    if (result.domain != null)
      return result;
  }

  return NO_ERROR;
}

error_t
map_container_values(
  map_container_iptr_t map,
  container_mptr_t     container
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  if (container == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_RECEIVER
    };

  malunal_size_t stride = 0;
  impl_iptr_t    self   = (impl_iptr_t)map;
  error_t        result = container_stride(container, &stride);
  if (result.domain != null)
    return result;

  if (stride != self->valstride)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_FAILURE
    };

  malunal_uint8_t* control = map_container_control(self);
  for (malunal_size_t index = 0; index < self->capacity; index++) {
    if (control[index] != MAP_CONTROL_FILLED)
      continue;

    result = container_append(container, map_container_value(self, index));
    if (result.domain != null)
      return result;
  }

  return NO_ERROR;
}

error_t
map_container_resize(
  map_container_mptr_t map,
  malunal_size_t       capacity
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  impl_mptr_t self = (impl_mptr_t)map;
  return map_container_realloc(self, capacity);
}

error_t
map_container_reserve(
  map_container_mptr_t map,
  malunal_size_t       capacity
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  impl_mptr_t self = (impl_mptr_t)map;
  return map_container_realloc(self, (capacity * 4) / 3 + 1);
}

error_t
map_container_append(
  map_container_mptr_t map,
  map_pair_iptr_t      pair
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  if (pair == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_FAILURE
    };

  return map_container_insert(map, pair->key, pair->value);
}

error_t
map_container_insert(
  map_container_mptr_t map,
  malunal_iptr_t       key,
  malunal_iptr_t       value
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  error_t     result = NO_ERROR;
  impl_mptr_t self   = (impl_mptr_t)map;
  if (self->capacity == 0) {
    result = map_container_realloc(self, MALUNAL_CONTAINERS_MAP_INIT_CAPACITY);
    if (result.domain != null)
      return result;
  }

  malunal_size_t slot;
  if (map_container_find(self, key, &slot)) {
    if (self->valstride != 0)
      memcpy(map_container_value(self, slot), value, self->valstride);
    return NO_ERROR;
  }

  if (slot == MAP_NO_SLOT ||
      (float)(self->count + 1) / self->capacity >= MAP_LOAD_FACTOR) {
    result = map_container_realloc(self, self->capacity * 2);
    if (result.domain != null)
      return result;

    map_container_find(self, key, &slot);
  }

  if (slot == MAP_NO_SLOT)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_FAILURE
    };

  map_container_place(self, slot, key, value);
  self->count++;
  return NO_ERROR;
}

error_t
map_container_remove(
  map_container_mptr_t map,
  malunal_iptr_t       key
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  malunal_size_t slot;
  impl_mptr_t    self = (impl_mptr_t)map;
  if (!map_container_find(self, key, &slot))
    return NO_ERROR;

  malunal_uint8_t* control = map_container_control(self);
  malunal_size_t   mask    = self->capacity - 1;
  control[slot] = MAP_CONTROL_ERASED;
  self->count--;

  if (control[(slot + 1) & mask] != MAP_CONTROL_EMPTY)
    return NO_ERROR;

  malunal_size_t index = slot;
  while (control[index] == MAP_CONTROL_ERASED) {
    control[index] = MAP_CONTROL_EMPTY;
    index = (index - 1) & mask;
  }

  return NO_ERROR;
}

error_t
map_container_contains(
  map_container_iptr_t map,
  malunal_iptr_t       key
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  malunal_size_t slot;
  impl_iptr_t    self = (impl_iptr_t)map;
  return map_container_find(self, key, &slot)
    ? NO_ERROR
    : (error_t) {
        .domain = &ERROR_DOMAIN_CONTAINER_T,
        .code   = CONTAINER_ERROR_FAILURE
      };
}

error_t
map_container_clear(
  map_container_mptr_t map
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  impl_mptr_t self = (impl_mptr_t)map;
  if (self->buffer != null)
    memset(self->buffer, MAP_CONTROL_EMPTY, self->capacity);

  self->count = 0;
  return NO_ERROR;
}

error_t
map_container_get(
  map_container_iptr_t map,
  malunal_iptr_t       key,
  malunal_mptr_t       value
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  malunal_size_t slot;
  impl_iptr_t    self = (impl_iptr_t)map;
  if (!map_container_find(self, key, &slot))
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_OUT_OF_BOUNDS
    };

  if (self->valstride != 0)
    memcpy(value, map_container_value(self, slot), self->valstride);
  return NO_ERROR;
}

error_t
map_container_set(
  map_container_mptr_t map,
  malunal_iptr_t       key,
  malunal_iptr_t       value
) {
  if (map == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_NULL_CONTAINER
    };

  malunal_size_t slot;
  impl_mptr_t    self = (impl_mptr_t)map;
  if (!map_container_find(self, key, &slot))
    return (error_t) {
      .domain = &ERROR_DOMAIN_CONTAINER_T,
      .code   = CONTAINER_ERROR_OUT_OF_BOUNDS
    };

  if (self->valstride != 0)
    memcpy(map_container_value(self, slot), value, self->valstride);
  return NO_ERROR;
}