#pragma once

#include "../../Core/Type/RuType_Base.h"
#include "../../Core/Memory/RuMemUtil.h"

#pragma managed(push, off)

// ************************************************************************************************************************************************************

template< class T > inline T Align( const T Ptr, INT Alignment )
{
	return (T)(((PTRVALUE)Ptr + Alignment - 1) & ~(Alignment-1));
}

// ************************************************************************************************************************************************************

class CRuMalloc_Windows
{
private:
	// Counts
	enum { POOL_COUNT = 42        };
	enum { POOL_MAX   = 32768 + 1 };

	// Forward declarations
	struct FreeMem;
	struct FPoolTable;

	// Memory pool info (32 bytes)
	struct PoolInfo
	{
		DWORD			m_bytes;			// Bytes allocated for pool
		DWORD			m_osBytes;			// Bytes aligned to page size
		DWORD			m_taken;			// Number of allocated elements in this pool, when counts down to zero can free the entire pool
		BYTE*			m_memBase;			// Memory base
		FPoolTable*		m_table;			// Index of pool
		FreeMem*		m_firstFreeMem;		// Pointer to first free memory in this pool
		PoolInfo*		m_next;
		PoolInfo**		m_prevLink;

		void Link(PoolInfo *&before)
		{
			if(before)
			{
				before->m_prevLink = &m_next;
			}

			m_next = before;
			m_prevLink = &before;
			before = this;
		}

		void Unlink()
		{
			if(m_next)
			{
				m_next->m_prevLink = m_prevLink;
			}

			*m_prevLink = m_next;
		}
	};

	// Information about a piece of free memory (8 bytes)
	struct FreeMem
	{
		FreeMem*		m_next;				// Next or MemLastPool[], always in order by pool
		DWORD			m_blocks;			// Number of consecutive free blocks here, at least 1

		PoolInfo*		GetPool()
		{
			return (PoolInfo *) ((INT) this & 0xFFFF0000);
		}
	};

	// Pool table
	struct FPoolTable
	{
		PoolInfo*		m_firstPool;
		PoolInfo*		m_exaustedPool;
		DWORD			m_blockSize;
	};

	// Variables
	INT32				m_pageSize;
	FPoolTable			m_poolTable[POOL_COUNT];
	FPoolTable			m_osTable;
	PoolInfo*			m_poolIndirect[256];
	FPoolTable*			m_memSizeToPoolTable[POOL_MAX];
	INT32				m_memInit;
	INT32				m_osCurrent;
	INT32				m_osPeak;
	INT32				m_usedCurrent;
	INT32				m_usedPeak;
	INT32				m_currentAllocs;
	INT32				m_totalAllocs;
	DOUBLE				m_memTime;

public:
	CRuMalloc_Windows()
	:	m_memInit(0),
		m_osCurrent(0),
		m_osPeak(0),
		m_usedCurrent(0),
		m_usedPeak(0),
		m_currentAllocs(0),
		m_totalAllocs(0),
		m_memTime(0.0)
	{
	}

	void *Allocate(DWORD size)
	{
		FreeMem *Free;
		if(size < POOL_MAX)
		{
			// Allocate from pool
			FPoolTable *Table = m_memSizeToPoolTable[size];
			PoolInfo *Pool = Table->m_firstPool;
			if(!Pool)
			{
				// Must create a new pool
				DWORD Blocks  = 65536 / Table->m_blockSize;
				DWORD Bytes   = Blocks * Table->m_blockSize;

				// Allocate memory
				Free = (FreeMem *) VirtualAlloc(NULL, Bytes, MEM_COMMIT, PAGE_READWRITE);
				if(!Free)
				{
					OutOfMemory();
				}

				// Create pool in the indirect table
				PoolInfo *&Indirect = m_poolIndirect[((DWORD) Free >> 24)];
				if(!Indirect)
				{
					Indirect = CreateIndirect();
				}

				Pool = &Indirect[((DWORD) Free >> 16) & 255];

				// Init pool
				Pool->Link(Table->m_firstPool);
				Pool->m_memBase = (BYTE *) Free;
				Pool->m_bytes = Bytes;
				Pool->m_osBytes = Align(Bytes, m_pageSize);
				Pool->m_table = Table;
				Pool->m_taken = 0;
				Pool->m_firstFreeMem = Free;

				// Create first free item
				Free->m_blocks = Blocks;
				Free->m_next = NULL;
			}

			// Pick first available block and unlink it
			Pool->m_taken++;
			Free = (FreeMem *) ((BYTE *) Pool->m_firstFreeMem + --Pool->m_firstFreeMem->m_blocks * Table->m_blockSize);
			if(Pool->m_firstFreeMem->m_blocks == 0)
			{
				Pool->m_firstFreeMem = Pool->m_firstFreeMem->m_next;
				if(!Pool->m_firstFreeMem)
				{
					// Move to exhausted list
					Pool->Unlink();
					Pool->Link(Table->m_exaustedPool);
				}
			}
		}
		else
		{
			// Use OS for large allocations
			INT32 AlignedSize = Align(size, m_pageSize);
			Free = (FreeMem *) VirtualAlloc(NULL, AlignedSize, MEM_COMMIT, PAGE_READWRITE);
			if(!Free)
			{
				OutOfMemory();
			}

			// Create indirect
			PoolInfo *&Indirect = m_poolIndirect[((DWORD) Free >> 24)];
			if(!Indirect)
			{
				Indirect = CreateIndirect();
			}

			// Init pool
			PoolInfo *Pool = &Indirect[((DWORD) Free >> 16) & 255];
			Pool->m_memBase = (BYTE *) Free;
			Pool->m_bytes = size;
			Pool->m_osBytes = AlignedSize;
			Pool->m_table = &m_osTable;
		}

		return Free;
	}

	void *Reallocate(void *Ptr, DWORD NewSize)
	{
		void *NewPtr = Ptr;
		if(Ptr && NewSize)
		{
			PoolInfo *Pool = &m_poolIndirect[(DWORD) Ptr >> 24][((DWORD) Ptr >> 16) & 255];
			if(Pool->m_table != &m_osTable)
			{
				// Allocated from pool, so grow or shrink if necessary
				if(NewSize>Pool->m_table->m_blockSize || m_memSizeToPoolTable[NewSize] != Pool->m_table)
				{
					NewPtr = Allocate(NewSize);
					RuMemcpy(NewPtr, Ptr, min(NewSize, Pool->m_table->m_blockSize));
					Deallocate(Ptr);
				}
			}
			else
			{
				// Allocated from OS
				if(NewSize>Pool->m_osBytes || NewSize * 3 < Pool->m_osBytes * 2)
				{
					// Grow or shrink
					NewPtr = Allocate(NewSize);
					RuMemcpy(NewPtr, Ptr, min(NewSize, Pool->m_bytes));
					Deallocate(Ptr);
				}
				else
				{
					// Keep as-is, reallocation isn't worth the overhead.
					Pool->m_bytes = NewSize;
				}
			}
		}
		else if(Ptr == NULL)
		{
			NewPtr = Allocate(NewSize);
		}
		else
		{
			Deallocate(Ptr);
			NewPtr = NULL;
		}

		return NewPtr;
	}

	void Deallocate(void *Ptr)
	{
		if(!Ptr)
		{
			return;
		}

		// Windows version
		PoolInfo *Pool = &m_poolIndirect[(DWORD) Ptr >> 24][((DWORD) Ptr >> 16) & 255];
		if(Pool->m_table != &m_osTable)
		{
			// If this pool was exhausted, move to available list
			if(!Pool->m_firstFreeMem)
			{
				Pool->Unlink();
				Pool->Link(Pool->m_table->m_firstPool);
			}

			// Free a pooled allocation
			FreeMem *Free = (FreeMem *) Ptr;
			Free->m_blocks = 1;
			Free->m_next = Pool->m_firstFreeMem;
			Pool->m_firstFreeMem = Free;

			// Free this pool
			if(--Pool->m_taken == 0)
			{
				// Free the OS memory
				Pool->Unlink();
				VirtualFree(Pool->m_memBase, 0, MEM_RELEASE);
			}
		}
		else
		{
			// Free an OS allocation
			VirtualFree(Ptr, 0, MEM_RELEASE);
		}
	}

	void Initialize()
	{
		m_memInit = 1;

		// Get OS page size
		SYSTEM_INFO SI;
		GetSystemInfo(&SI);
		m_pageSize = SI.dwPageSize;

		// Init tables
		m_osTable.m_firstPool = NULL;
		m_osTable.m_exaustedPool = NULL;
		m_osTable.m_blockSize = 0;

		m_poolTable[0].m_firstPool = NULL;
		m_poolTable[0].m_exaustedPool = NULL;
		m_poolTable[0].m_blockSize = 8;

		for(DWORD i = 1; i < 5; i++)
		{
			m_poolTable[i].m_firstPool = NULL;
			m_poolTable[i].m_exaustedPool = NULL;
			m_poolTable[i].m_blockSize = (8 << ((i + 1) >> 2)) + (2 << i);
		}

		for(DWORD i = 5; i < POOL_COUNT; i++)
		{
			m_poolTable[i].m_firstPool = NULL;
			m_poolTable[i].m_exaustedPool = NULL;
			m_poolTable[i].m_blockSize = (4 + ((i + 7) & 3)) << (1 + ((i + 7) >> 2));
		}

		for(DWORD i = 0; i < POOL_MAX; i++)
		{
			DWORD Index;
			for(Index = 0; m_poolTable[Index].m_blockSize < i; Index++);
			m_memSizeToPoolTable[i] = &m_poolTable[Index];
		}

		for(DWORD i = 0; i < 256; i++)
		{
			m_poolIndirect[i] = NULL;
		}
	}

protected:
	// Implementation
	void OutOfMemory()
	{
	}

	PoolInfo *CreateIndirect()
	{
		PoolInfo *indirect = (PoolInfo *) VirtualAlloc(NULL, 256 * sizeof(PoolInfo), MEM_COMMIT, PAGE_READWRITE);

		if(!indirect)
		{
			OutOfMemory();
		}

		return indirect;
	}
};

// ************************************************************************************************************************************************************

#pragma managed(pop)
