sourcemod/knight/shared/KeCodeAllocator.cpp

#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;
}