#include <ZUtil/ZAllocWindow.hpp>
#include <ZUtil/ZAssert.hpp>

#include <SST/SST_OS.h>

#define ALLOC_BITS		(0xF00D)	// Used to flag a section of memory as 'allocated'
#define DEALLOC_BITS	(0xDEAD)	// Used to flag a section of memory as 'deallocated'
#define LOOP_BITS		(0xCAB0)	// Used to flag that the next allocated section is back at the beginning of the buffer

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

ZAllocWindow::ZAllocWindow()
: Buffer(NULL), Begin(NULL), End(NULL),
  BufferSize(0), bOwnsBuffer(true),
  NumAllocs(0)
{

}

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

ZAllocWindow::ZAllocWindow( size_t _buffersize )
: Buffer(NULL), Begin(NULL), End(NULL),
  BufferSize(_buffersize), bOwnsBuffer(true),
  NumAllocs(0)
{
	Resize(_buffersize);
}

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

ZAllocWindow::~ZAllocWindow()
{
	if (bOwnsBuffer)
		free(Buffer);
}

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

void* ZAllocWindow::Allocate( size_t _size )
{
	// boundary is the location in memory that we cannot write past
	uintptr_t boundary;

	// determine where our boundary is
	uintptr_t pool_begin = (uintptr_t)Buffer;
	uintptr_t pool_end	 = pool_begin + BufferSize;
	uintptr_t ptr_begin  = (uintptr_t)Begin;
	uintptr_t ptr_end	 = (uintptr_t)End;

	boundary = (ptr_begin >  ptr_end ? ptr_begin : 
		        ptr_begin <  ptr_end ? pool_end  :
									 ( NumAllocs > 0 ? ptr_end : 
													   pool_end ) );

	// align segment to word boundary
	uintptr_t seg_begin = (uintptr_t)End;
	uintptr_t seg_end   = (uintptr_t)SST_OS_AlignPtr((void*)(seg_begin + 2*sizeof(uint32_t) + _size), sizeof(void*));
	
	// check segment against boundary
	if (seg_end > boundary) {

		// if the segment has not just run up against the pool end, we are done
		if (boundary != pool_end) {
			return NULL;
		}

		// signal a loop (if less than 'end') and check again
		if (seg_begin < pool_end) {
			uint32_t* ptr_32 = (uint32_t*)seg_begin;
			ptr_32[0] = (uint32_t)LOOP_BITS;
		}

		boundary = ptr_begin;

		seg_begin = pool_begin;
		seg_end   = (uintptr_t)SST_OS_AlignPtr((void*)(seg_begin +  2 * sizeof(uint32_t) + _size), sizeof(void*));
	}

	// check segment against boundary again
	if (seg_end <= boundary) {
		NumAllocs++;

		// set size as first element in block
		uint32_t* ptr_32 = (uint32_t*)seg_begin;
		ptr_32[0] = (uint32_t)ALLOC_BITS;
		ptr_32[1] = (uint32_t)(seg_end - seg_begin);
				
		// set new End pointer to ptr_end
		End = (void*)seg_end;

		// if End == BufferEnd, End == 0
		if (End == (void*)((uintptr_t)Buffer + (uintptr_t)BufferSize))
			End = Buffer;

		// return the address of the first non-reserved element of the segment
		return &ptr_32[2];
	}

	return NULL;
}

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

void ZAllocWindow::Deallocate( void* _data )
{
	// check size before data
	uint32_t* ptr_32 = (uint32_t*)((uintptr_t)_data - (2*sizeof(uint32_t)));
	uint32_t  flag = ptr_32[0];
	uint32_t  size = ptr_32[1];

	// quick sanity check
	ZASSERT(flag == ALLOC_BITS, "ZAllocWindow passed invalid pointer to deallocate!");
	ZASSERT(size <  BufferSize, "ZAllocWindow passed corrupted pointer to deallocate!");
	
	// mark size slot as 'freed'
	ptr_32[0] = (uint32_t)DEALLOC_BITS;
	NumAllocs--;

	// see if this was the oldest allocated block, and start rolling up
	if (ptr_32 == Begin) {
		while (WHILE_TRUE) {
			ptr_32 = (uint32_t*)Begin;

			// if begin is at or past end, loop to the beginning
			if ((uintptr_t)ptr_32 >= ((uintptr_t)Begin + BufferSize)) {
				Begin = Buffer;
				continue;
			}

			// determine flag (size if needed)
			flag = ptr_32[0];

			// if this is a loop flag, start again at the beginning
			if (flag == LOOP_BITS) {
				Begin = Buffer;
				continue;
			}

			// if this is a dealloc flag, dealloc the segment, increment Begin
			if (flag == DEALLOC_BITS) {
				size = ptr_32[1];

				uintptr_t seg_start = (uintptr_t)ptr_32;
				uintptr_t seg_end   = seg_start + size;

				Begin = (void*)seg_end;

				#ifdef _DEBUG
				ptr_32[1] = DEALLOC_BITS;
				#endif

				continue;
			}

			break;
		}
	}
}

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

void ZAllocWindow::Reset()
{
	NumAllocs = 0;
	Begin = Buffer;
	End = Buffer;
}

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

bool ZAllocWindow::Resize( size_t _newsize )
{
	// don't resize if we have allocations or don't own the buffer
	if (NumAllocs > 0 || !bOwnsBuffer) {
		return false;
	}

	// check newsize is zero to free everything
	if (_newsize == 0) {
		if (Buffer != NULL) {
			free(Buffer);
			Buffer = NULL;
		}
		return true;
	}

	// free current buffer if not NULL
	if (Buffer != NULL) {
		free(Buffer);
	}

	// allocate new buffer and return success / failure
	Buffer = malloc(_newsize);
	Begin = Buffer;
	End = Buffer; 
	BufferSize = _newsize;

	return Buffer != NULL;
}

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

bool ZAllocWindow::SetBuffer( void* _buffer, size_t _buffersize )
{
	// don't accept buffer if we have allocations
	if (NumAllocs > 0) {
		return 0;
	}

	// if the new buffer is null, just reallocate to buffersize
	if (_buffer == NULL) {
		bOwnsBuffer = true;
		return Resize(_buffersize);
	}

	// Get rid of the current buffer if we need to
	if (Buffer != NULL) {
		free(Buffer);
	}

	// lose ownership and set parameters, return success
	bOwnsBuffer = false;
	Buffer = _buffer;
	Begin = Buffer;
	End = Buffer; 
	BufferSize = _buffersize;

	return true;
}