#include <KePlatform.h>

 #include <assert.h>

 #include <string.h>

 #if defined KE_PLATFORM_WINDOWS

 #include <windows.h>

 #elif defined KE_PLATFORM_POSIX

 #include <unistd.h>

 #include <stdlib.h>

 #include <sys/mman.h>

 #else

 #error "TODO"

 #endif

 #include "KeCodeAllocator.h"

 #define ALIGNMENT	16

 using namespace Knight;

 struct KeFreedCode;

 /**

  * Defines a region of memory that is made of pages. 

  */

 struct KeCodeRegion

 {

 	KeCodeRegion *next;

 	unsigned char *block_start;

 	unsigned char *block_pos;

 	KeFreedCode *free_list;

 	size_t total_size;

 	size_t end_free;

 	size_t total_free;

 };

 /**

  * Defines freed code.  We keep the size here because 

  * when we touch the linked list we don't want to dirty pages.

  */

 struct KeFreedCode

 {

 	KeCodeRegion *region;

 	unsigned char *block_start;

 	size_t size;

 	KeFreedCode *next;

 };

 struct KeSecret

 {

 	KeCodeRegion *region;

 	size_t size;

 };

 class Knight::KeCodeCache

 {

 public:

 	/**

 	 * First region that is live for use.

 	 */

 	KeCodeRegion *first_live;

 	/**

 	 * First region that is full but has free entries.

 	 */

 	KeCodeRegion *first_partial;

 	/**

 	 * First region that is full.

 	 */

 	KeCodeRegion *first_full;

 	/**

 	 * Page granularity and size.

 	 */

 	unsigned int page_size;

 	unsigned int page_granularity;

 	/**

 	 * This isn't actually for code, this is the node cache.

 	 */

 	KeCodeRegion *node_cache;

 	KeFreedCode *free_node_list;

 };

 KeCodeCache *Knight::KE_CreateCodeCache()

 {

 	KeCodeCache *cache;

 	cache = new KeCodeCache;

 	memset(cache, 0, sizeof(KeCodeCache));

 #if defined KE_PLATFORM_WINDOWS

 	SYSTEM_INFO info;

 	GetSystemInfo(&info);

 	cache->page_size = info.dwPageSize;

 	cache->page_granularity = info.dwAllocationGranularity;

 #else

 	cache->page_size = cache->page_granularity = sysconf(_SC_PAGESIZE);

 #endif

 	return cache;

 }

 inline size_t MinAllocSize()

 {

 	size_t size;

 	size = sizeof(KeSecret);

 	size += ALIGNMENT;

 	size -= size % ALIGNMENT;

 	return size;

 }

 inline size_t ke_GetAllocSize(size_t size)

 {

 	size += sizeof(KeSecret);

 	size += ALIGNMENT;

 	size -= size % ALIGNMENT;

 	return size;

 }

 void *ke_AllocInRegion(KeCodeCache *cache,

 					   KeCodeRegion **prev,

 					   KeCodeRegion *region,

 					   unsigned char *ptr,

 					   size_t alloc_size,

 					   bool is_live)

 {

 	KeSecret *secret;

 	/* Squirrel some info in the alloc. */

 	secret = (KeSecret *)ptr;

 	secret->region = region;

 	secret->size = alloc_size;

 	ptr += sizeof(KeSecret);

 	region->total_free -= alloc_size;

 	/* Check if we can't use the fast-path anymore. */

 	if ((is_live && region->end_free < MinAllocSize())

 		|| (!is_live && region->total_free < MinAllocSize()))

 	{

 		KeCodeRegion **start;

 		*prev = region->next;

 		/* Select the appropriate arena. */

 		if (is_live)

 		{

 			if (region->total_free < MinAllocSize())

 			{

 				start = &cache->first_full;

 			}

 			else

 			{

 				start = &cache->first_partial;

 			}

 		}

 		else

 		{

 			start = &cache->first_full;

 		}

 		region->next = *start;

 		*start = region;

 	}

 	return ptr;

 }

 void *ke_AllocFromLive(KeCodeCache *cache, size_t size)

 {

 	void *ptr;

 	size_t alloc_size;

 	KeCodeRegion *region, **prev;

 	region = cache->first_live;

 	prev = &cache->first_live;

 	alloc_size = ke_GetAllocSize(size);

 	while (region != NULL)

 	{

 		if (region->end_free >= alloc_size)

 		{

 			/* Yay! We can do a simple alloc here. */

 			ptr = ke_AllocInRegion(cache, prev, region, region->block_pos, alloc_size, true);

 			/* Update our counters. */

 			region->block_pos += alloc_size;

 			region->end_free -= alloc_size;

 			return ptr;

 		}

 		prev = &region->next;

 		region = region->next;

 	}

 	return NULL;

 }

 void *ke_AllocFromPartial(KeCodeCache *cache, size_t size)

 {

 	void *ptr;

 	size_t alloc_size;

 	KeCodeRegion *region, **prev;

 	region = cache->first_partial;

 	prev = &cache->first_partial;

 	alloc_size = ke_GetAllocSize(size);

 	while (region != NULL)

 	{

 		if (region->total_free >= alloc_size)

 		{

 			KeFreedCode *node, **last;

 			assert(region->free_list != NULL);

 			last = &region->free_list;

 			node = region->free_list;

 			while (node != NULL)

 			{

 				if (node->size >= alloc_size)

 				{

 					/* Use this node */

 					ptr = ke_AllocInRegion(cache, prev, region, node->block_start, alloc_size, false);

 					region->total_free -= node->size;

 					*last = node->next;

 					/* Make sure bookkeepping is correct. */

 					assert((region->free_list == NULL && region->total_free == 0)

 						   || (region->free_list != NULL && region->total_free != 0));

 					/* Link us back into the free node list. */

 					node->next = cache->free_node_list;

 					cache->free_node_list = node->next;

 					return ptr;

 				}

 				last = &node->next;

 				node = node->next;

 			}

 		}

 		prev = &region->next;

 		region = region->next;

 	}

 	return NULL;

 }

 KeCodeRegion *ke_AddRegionForSize(KeCodeCache *cache, size_t size)

 {

 	KeCodeRegion *region;

 	region = new KeCodeRegion;

 	size = ke_GetAllocSize(size);

 	size += cache->page_granularity * 2;

 	size -= size % cache->page_granularity;

 #if defined KE_PLATFORM_WINDOWS

 	region->block_start = (unsigned char *)VirtualAlloc(NULL, size, MEM_COMMIT|MEM_RESERVE, PAGE_EXECUTE_READWRITE);

 #elif defined KE_PLATFORM_POSIX

 	region->block_start = (unsigned char *)valloc(size);

 	if (mprotect(region->block_start, size, PROT_READ|PROT_WRITE|PROT_EXEC) == -1)

 	{

 		free(region->block_start);

 		delete region;

 		return NULL;

 	}

 #else

 #error "TODO"

 #endif

 	if (region->block_start == NULL)

 	{

 		delete region;

 		return NULL;

 	}

 	region->block_pos = region->block_start;

 	region->end_free = region->total_free = region->total_size = size;

 	region->next = cache->first_live;

 	cache->first_live = region;

 	region->free_list = NULL;

 	return region;

 }

 void *Knight::KE_AllocCode(KeCodeCache *cache, size_t size)

 {

 	void *ptr;

 	/* Check live easy-adds */

 	if (cache->first_live != NULL)

 	{

 		if ((ptr = ke_AllocFromLive(cache, size)) != NULL)

 		{

 			return ptr;

 		}

 	}

 	/* Try looking in the free lists */

 	if (cache->first_partial != NULL)

 	{

 		if ((ptr = ke_AllocFromPartial(cache, size)) != NULL)

 		{

 			return ptr;

 		}

 	}

 	/* Create a new region */

 	if (ke_AddRegionForSize(cache, size) == NULL)

 	{

 		return NULL;

 	}

 	return ke_AllocFromLive(cache, size);

 }

 KeFreedCode *ke_GetFreeNode(KeCodeCache *cache)

 {

 	KeFreedCode *ret;

 	if (cache->free_node_list != NULL)

 	{

 		ret = cache->free_node_list;

 		cache->free_node_list = ret->next;

 		return ret;

 	}

 	/* See if the current free node region has space. */

 	if (cache->node_cache != NULL

 		&& cache->node_cache->end_free >= sizeof(KeFreedCode))

 	{

 		ret = (KeFreedCode *)cache->node_cache->block_pos;

 		cache->node_cache->block_pos += sizeof(KeFreedCode);

 		cache->node_cache->total_free -= sizeof(KeFreedCode);

 		cache->node_cache->end_free -= sizeof(KeFreedCode);

 		return ret;

 	}

 	/* Otherwise, we need to alloc a new region. */

 	KeCodeRegion *region = new KeCodeRegion;

 	region->block_start = new unsigned char[cache->page_size / sizeof(KeFreedCode)];

 	region->block_pos = region->block_start + sizeof(KeFreedCode);

 	region->total_size = cache->page_size / sizeof(KeFreedCode);

 	region->total_free = region->end_free = (region->total_size - sizeof(KeFreedCode));

 	region->free_list = NULL;

 	region->next = cache->node_cache;

 	cache->node_cache = region;

 	return (KeFreedCode *)region->block_start;

 }

 void Knight::KE_FreeCode(KeCodeCache *cache, void *code)

 {

 	KeSecret *secret;

 	KeFreedCode *node;

 	unsigned char *ptr;

 	KeCodeRegion *region;

 	ptr = (unsigned char *)code;

 	secret = (KeSecret *)(ptr - sizeof(KeSecret));

 	region = secret->region;

 	node = ke_GetFreeNode(cache);

 	node->block_start = (unsigned char *)code;

 	node->next = region->free_list;

 	region->free_list = node;

 	node->region = region;

 	node->size = secret->size;

 }

 KeCodeRegion *ke_DestroyRegion(KeCodeRegion *region)

 {

 	KeCodeRegion *next;

 	next = region->next;

 #if defined KE_PLATFORM_WINDOWS

 	VirtualFree(region->block_start, 0, MEM_RELEASE);

 #else

 	free(region->block_start);

 #endif

 	delete region;

 	return next;

 }

 void ke_DestroyRegionChain(KeCodeRegion *first)

 {

 	while (first != NULL)

 	{

 		first = ke_DestroyRegion(first);

 	}

 }

 void Knight::KE_DestroyCodeCache(KeCodeCache *cache)

 {

 	/* Destroy every region and call it a day. */

 	ke_DestroyRegionChain(cache->first_full);

 	ke_DestroyRegionChain(cache->first_live);

 	ke_DestroyRegionChain(cache->first_partial);

 	/* We use normal malloc for node cache regions */

 	KeCodeRegion *region, *next;

 	region = cache->node_cache;

 	while (region != NULL)

 	{

 		next = region->next;

 		delete [] region->block_start;

 		delete region;

 		region = next;

 	}

 	delete cache;

 }