// Copyright Epic Games, Inc. All Rights Reserved.

#include <zencore/memory/mallocansi.h>

#include <zencore/intmath.h>
#include <zencore/memory/align.h>
#include <zencore/windows.h>

#if ZEN_PLATFORM_LINUX
#	define PLATFORM_USE_ANSI_POSIX_MALLOC 1
#endif

#if ZEN_PLATFORM_MAC
#	define PLATFORM_USE_CUSTOM_MEMALIGN 1
#endif

#ifndef PLATFORM_USE_ANSI_MEMALIGN
#	define PLATFORM_USE_ANSI_MEMALIGN 0
#endif

#ifndef PLATFORM_USE_ANSI_POSIX_MALLOC
#	define PLATFORM_USE_ANSI_POSIX_MALLOC 0
#endif

#ifndef PLATFORM_USE_CUSTOM_MEMALIGN
#	define PLATFORM_USE_CUSTOM_MEMALIGN 0
#endif

#if PLATFORM_USE_ANSI_POSIX_MALLOC
#	include <malloc.h>
#	include <string.h>
#endif

#define MALLOC_ANSI_USES__ALIGNED_MALLOC ZEN_PLATFORM_WINDOWS

namespace zen {

//////////////////////////////////////////////////////////////////////////

void*
AnsiMalloc(size_t Size, uint32_t Alignment)
{
#if MALLOC_ANSI_USES__ALIGNED_MALLOC
	void* Result = _aligned_malloc(Size, Alignment);
#elif PLATFORM_USE_ANSI_POSIX_MALLOC
	void* Result;
	if (posix_memalign(&Result, Alignment, Size) != 0)
	{
		Result = nullptr;
	}
#elif PLATFORM_USE_ANSI_MEMALIGN
	Result = reallocalign(Ptr, NewSize, Alignment);
#elif PLATFORM_USE_CUSTOM_MEMALIGN
	void* Ptr	 = malloc(Size + Alignment + sizeof(void*) + sizeof(size_t));
	void* Result = nullptr;
	if (Ptr)
	{
		Result															= Align((uint8_t*)Ptr + sizeof(void*) + sizeof(size_t), Alignment);
		*((void**)((uint8_t*)Result - sizeof(void*)))					= Ptr;
		*((size_t*)((uint8_t*)Result - sizeof(void*) - sizeof(size_t))) = Size;
	}
#else
#	error Unknown allocation path
#endif

	return Result;
}

size_t
AnsiGetAllocationSize(void* Original)
{
#if MALLOC_ANSI_USES__ALIGNED_MALLOC
	return _aligned_msize(Original, 16, 0);	 // TODO: incorrectly assumes alignment of 16
#elif PLATFORM_USE_ANSI_POSIX_MALLOC || PLATFORM_USE_ANSI_MEMALIGN
	return malloc_usable_size(Original);
#elif PLATFORM_USE_CUSTOM_MEMALIGN
	return *((size_t*)((uint8_t*)Original - sizeof(void*) - sizeof(size_t)));
#else
#	error Unknown allocation path
#endif
}

void*
AnsiRealloc(void* Ptr, size_t NewSize, uint32_t Alignment)
{
	void* Result = nullptr;

#if MALLOC_ANSI_USES__ALIGNED_MALLOC
	if (Ptr && NewSize)
	{
		Result = _aligned_realloc(Ptr, NewSize, Alignment);
	}
	else if (Ptr == nullptr)
	{
		Result = _aligned_malloc(NewSize, Alignment);
	}
	else
	{
		_aligned_free(Ptr);
		Result = nullptr;
	}
#elif PLATFORM_USE_ANSI_POSIX_MALLOC
	if (Ptr && NewSize)
	{
		size_t UsableSize = malloc_usable_size(Ptr);
		if (posix_memalign(&Result, Alignment, NewSize) != 0)
		{
			Result = nullptr;
		}
		else if (UsableSize)
		{
			memcpy(Result, Ptr, Min(NewSize, UsableSize));
		}
		free(Ptr);
	}
	else if (Ptr == nullptr)
	{
		if (posix_memalign(&Result, Alignment, NewSize) != 0)
		{
			Result = nullptr;
		}
	}
	else
	{
		free(Ptr);
		Result = nullptr;
	}
#elif PLATFORM_USE_CUSTOM_MEMALIGN
	if (Ptr && NewSize)
	{
		// Can't use realloc as it might screw with alignment.
		Result		   = AnsiMalloc(NewSize, Alignment);
		size_t PtrSize = AnsiGetAllocationSize(Ptr);
		memcpy(Result, Ptr, Min(NewSize, PtrSize));
		AnsiFree(Ptr);
	}
	else if (Ptr == nullptr)
	{
		Result = AnsiMalloc(NewSize, Alignment);
	}
	else
	{
		free(*((void**)((uint8_t*)Ptr - sizeof(void*))));
		Result = nullptr;
	}
#else
#	error Unknown allocation path
#endif

	return Result;
}

void
AnsiFree(void* Ptr)
{
#if MALLOC_ANSI_USES__ALIGNED_MALLOC
	_aligned_free(Ptr);
#elif PLATFORM_USE_ANSI_POSIX_MALLOC || PLATFORM_USE_ANSI_MEMALIGN
	free(Ptr);
#elif PLATFORM_USE_CUSTOM_MEMALIGN
	if (Ptr)
	{
		free(*((void**)((uint8_t*)Ptr - sizeof(void*))));
	}
#else
#	error Unknown allocation path
#endif
}

//////////////////////////////////////////////////////////////////////////

FMallocAnsi::FMallocAnsi()
{
#if ZEN_PLATFORM_WINDOWS
	// Enable low fragmentation heap - http://msdn2.microsoft.com/en-US/library/aa366750.aspx
	intptr_t CrtHeapHandle = _get_heap_handle();
	ULONG	 EnableLFH	   = 2;
	HeapSetInformation((void*)CrtHeapHandle, HeapCompatibilityInformation, &EnableLFH, sizeof(EnableLFH));
#endif
}

void*
FMallocAnsi::TryMalloc(size_t Size, uint32_t Alignment)
{
	Alignment = Max(Size >= 16 ? (uint32_t)16 : (uint32_t)8, Alignment);

	void* Result = AnsiMalloc(Size, Alignment);

	return Result;
}

void*
FMallocAnsi::Malloc(size_t Size, uint32_t Alignment)
{
	void* Result = TryMalloc(Size, Alignment);

	if (Result == nullptr && Size)
	{
		OutOfMemory(Size, Alignment);
	}

	return Result;
}

void*
FMallocAnsi::TryRealloc(void* Ptr, size_t NewSize, uint32_t Alignment)
{
	Alignment = Max(NewSize >= 16 ? (uint32_t)16 : (uint32_t)8, Alignment);

	void* Result = AnsiRealloc(Ptr, NewSize, Alignment);

	return Result;
}

void*
FMallocAnsi::Realloc(void* Ptr, size_t NewSize, uint32_t Alignment)
{
	void* Result = TryRealloc(Ptr, NewSize, Alignment);

	if (Result == nullptr && NewSize != 0)
	{
		OutOfMemory(NewSize, Alignment);
	}

	return Result;
}

void
FMallocAnsi::Free(void* Ptr)
{
	AnsiFree(Ptr);
}

bool
FMallocAnsi::GetAllocationSize(void* Original, size_t& SizeOut)
{
	if (!Original)
	{
		return false;
	}

#if MALLOC_ANSI_USES__ALIGNED_MALLOC
	ZEN_UNUSED(SizeOut);
	return false;
#else
	SizeOut = AnsiGetAllocationSize(Original);
	return true;
#endif
}

}  // namespace zen