#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <nllibc.h>

#define ALIGN_SIZE 16

/* #define MEMORY_USE_MMAP */

#ifndef MEMORY_BACKWARD_SEARCH_DISABLE
#ifndef MEMORY_BACKWARD_SEARCH_NUM
#define MEMORY_BACKWARD_SEARCH_NUM 64
#endif
#endif

#ifndef MEMORY_AREA_CACHE_DISABLE
#ifndef MEMORY_AREA_CACHE_NUM
#define MEMORY_AREA_CACHE_NUM 256
#endif
#ifndef MEMORY_AREA_CACHE_MAXLEN
#define MEMORY_AREA_CACHE_MAXLEN 64
#endif
#endif

#ifndef MEMORY_AREA_HEADER_MAGIC_NUMBER_DISABLE
#ifndef MEMORY_AREA_HEADER_MAGIC_NUMBER
#define MEMORY_AREA_HEADER_MAGIC_NUMBER 0x12345678
#endif
#endif

struct area_header {
  struct area_header *next;
  struct area_header *prev;
#ifdef MEMORY_AREA_CACHE_NUM
  struct area_header *cache;
#endif
  int size;
#ifdef MEMORY_AREA_HEADER_MAGIC_NUMBER
  int magic;
#endif
};

static struct area_header *alloc_start = NULL;
static struct area_header *alloc_end = NULL;

#ifdef MEMORY_AREA_CACHE_NUM
struct area_cache {
  int length;
  struct area_header *area;
};

static struct area_cache area_cache[MEMORY_AREA_CACHE_NUM];
#endif

static int alignment(int size)
{
  return ((size + (ALIGN_SIZE - 1)) / ALIGN_SIZE) * ALIGN_SIZE;
}

static char *extend_memory(void)
{
  static char *p = NULL;
  if (p != NULL) {
    if (sbrk(0x1000) == (void *)-1)
      return NULL;
  }
  p = sbrk(0);
  if (p == (void *)-1)
    p = NULL;
  return p;
}

static int extend(void)
{
  struct area_header *new_alloc_end;

  new_alloc_end = (struct area_header *)extend_memory();
  if (new_alloc_end == NULL)
    return -1;

  new_alloc_end--;

  alloc_end->prev->next = new_alloc_end;

  memset(new_alloc_end, 0, sizeof(*new_alloc_end));
  new_alloc_end->prev = alloc_end->prev;
  new_alloc_end->next = NULL;

  alloc_end = new_alloc_end;

  return 0;
}

#ifdef MEMORY_USE_MMAP
static int expand(size_t size)
{
  struct area_header *new_alloc_start;
  struct area_header *new_alloc_end;
  int fd = -1;
  char *p;

  size += sizeof(*new_alloc_start) + sizeof(*new_alloc_end);
  size = (size + 0x1000) / 0x1000 + 0x1000;

  fd = open("/dev/zero", O_RDWR);
  if (fd < 0)
    goto err;

  p = (char *)mmap(NULL, size, PROT_READ|PROT_WRITE, MAP_PRIVATE, fd, 0);
  if (p == NULL)
    goto err;

  new_alloc_start = (struct area_header *)p;
  new_alloc_end = (struct area_header *)(p + size);

  new_alloc_end--;

  alloc_end->next = new_alloc_start;
  alloc_end->size = -1; /* for skip */

  memset(new_alloc_start, 0, sizeof(*new_alloc_start));
  memset(new_alloc_end, 0, sizeof(*new_alloc_end));
  new_alloc_start->prev = alloc_end;
  new_alloc_start->next = new_alloc_end;
  new_alloc_end->prev = new_alloc_start;
  new_alloc_end->next = NULL;

  alloc_end = new_alloc_end;

  return 0;

err:
  if (fd >= 0)
    close(fd);

  return -1;
}
#endif

static int extend_or_expand(size_t size)
{
  int r = -1;
  static int continuous = 1;

  if (continuous)
    r = extend();
#ifdef MEMORY_USE_MMAP
  if (r < 0) {
    r = expand(size);
    continuous = 0;
  }
#endif

  return r;
}

static int init(void)
{
  if (alloc_start == NULL) {
    alloc_start = (struct area_header *)extend_memory();
    alloc_end = (struct area_header *)extend_memory();

    if ((alloc_start == NULL) || (alloc_end == NULL))
      return -1;

    alloc_end--;

    memset(alloc_start, 0, sizeof(*alloc_start));
    memset(alloc_end, 0, sizeof(*alloc_end));
    alloc_start->prev = NULL;
    alloc_start->next = alloc_end;
    alloc_end->prev = alloc_start;
    alloc_end->next = NULL;

#ifdef MEMORY_AREA_CACHE_NUM
    memset(area_cache, 0, sizeof(area_cache));
#endif
  }
  return 0;
}

void *malloc(size_t size)
{
  struct area_header *area, *following;
  int aligned_size;
#ifdef MEMORY_AREA_CACHE_NUM
  int i;
#endif
#ifdef MEMORY_BACKWARD_SEARCH_NUM
  int n;
#endif

  if (init() < 0)
    return NULL;

  aligned_size = alignment(size);

#ifdef MEMORY_AREA_CACHE_NUM
  i = aligned_size / ALIGN_SIZE;
  if (i < MEMORY_AREA_CACHE_NUM) {
    area = area_cache[i].area;
    if (area != NULL) {
      area_cache[i].length--;
      area_cache[i].area = area->cache;
      area->cache = NULL;
      return area + 1;
    }
  }
#endif

#ifdef MEMORY_BACKWARD_SEARCH_NUM
  area = alloc_end;
  for (n = 0; n < MEMORY_BACKWARD_SEARCH_NUM; n++) {
    if (area->prev == NULL)
      break;
    area = area->prev;
  }
#else
  area = alloc_start;
#endif

  while (1) {
    for (; area->next; area = area->next) {
      if (area->size < 0) /* skip */
	continue;
      following = (struct area_header *)((char *)(area + 1) + area->size);
      if ((char *)area->next >= (char *)(following + 1) + aligned_size) {
	memset(following, 0, sizeof(*following));
	following->size = aligned_size;
#ifdef MEMORY_AREA_HEADER_MAGIC_NUMBER
	following->magic = MEMORY_AREA_HEADER_MAGIC_NUMBER;
#endif
	following->prev = area;
	following->next = area->next;
	area->next->prev = following;
	area->next = following;
	return following + 1;
      }
    }

    area = area->prev;

    if (extend_or_expand(size) < 0)
      break;
  }

  return NULL;
}

void *calloc(size_t number, size_t size)
{
  int s = size * number;
  void *p;

  p = malloc(s);
  if (p != NULL) {
    memset(p, 0, s);
  }

  return p;
}

void *realloc(void *ptr, size_t size)
{
  struct area_header *area = NULL;
  int aligned_size;
  void *p;

  if (init() < 0)
    return NULL;

  aligned_size = alignment(size);

  if (ptr != NULL)
    area = ((struct area_header *)ptr) - 1;

  if (area != NULL) {
    if (aligned_size <= area->size)
      return ptr;

    if ((char *)area->next >= (char *)(area + 1) + aligned_size) {
      area->size = aligned_size;
      return ptr;
    }
  }

  p = malloc(size);
  if (p == NULL)
    return NULL;

  if (area != NULL) {
    memcpy(p, ptr, area->size);
    free(ptr);
  }

  return p;
}

void free(void *ptr)
{
  struct area_header *area;
#ifdef MEMORY_AREA_CACHE_NUM
  int i;
#endif

  if (ptr == NULL)
    return;

  area = ((struct area_header *)ptr) - 1;

#ifdef MEMORY_AREA_HEADER_MAGIC_NUMBER
  if (area->magic != MEMORY_AREA_HEADER_MAGIC_NUMBER) {
    ;
  }
#endif

#ifdef MEMORY_AREA_CACHE_NUM
  i = area->size / ALIGN_SIZE;
  if ((i < MEMORY_AREA_CACHE_NUM)
#ifdef MEMORY_AREA_CACHE_MAXLEN
      && ((MEMORY_AREA_CACHE_MAXLEN < 0) ||
	  (area_cache[i].length < MEMORY_AREA_CACHE_MAXLEN))
#endif
      ) {
    area->cache = area_cache[i].area;
    area_cache[i].area = area;
    area_cache[i].length++;
    return;
  }
#endif

  area->prev->next = area->next;
  area->next->prev = area->prev;

  memset(area, 0, sizeof(*area));

  return;
}

int __nllibc_memory_stat(int *alloc_sizep, int *alloc_countp, int *free_sizep)
{
  struct area_header *area, *following;
  int alloc_size = 0, alloc_count = 0, free_size = 0;

  if (init() < 0)
    return -1;

  area = alloc_start;

  for (; area->next; area = area->next) {
    if (area->size < 0) /* skip */
      continue;
    alloc_size += area->size;
    alloc_count++;
    following = (struct area_header *)((char *)(area + 1) + area->size);
    free_size += (char *)area->next - (char *)following;
  }

  if (alloc_sizep)  *alloc_sizep  = alloc_size;
  if (alloc_countp) *alloc_countp = alloc_count;
  if (free_sizep)   *free_sizep   = free_size;

  return 0;
}
