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

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

#if ZEN_RPMALLOC_ENABLED

#	include "rpmalloc.h"
#	include <cstring>

/** Value we fill a memory block with after it is free, in UE_BUILD_DEBUG **/
#	define DEBUG_FILL_FREED (0xdd)

/** Value we fill a new memory block with, in UE_BUILD_DEBUG **/
#	define DEBUG_FILL_NEW (0xcd)

#	define ZEN_ENABLE_DEBUG_FILL 1

namespace zen {

FMallocRpmalloc::FMallocRpmalloc()
{
	rpmalloc_initialize(nullptr);
}

FMallocRpmalloc::~FMallocRpmalloc()
{
	rpmalloc_finalize();
}

void*
FMallocRpmalloc::TryMalloc(size_t Size, uint32_t Alignment)
{
	void* NewPtr = nullptr;

	if (Alignment != DEFAULT_ALIGNMENT)
	{
		Alignment = Max(uint32_t(Size >= 16 ? 16 : 8), Alignment);
		NewPtr	  = rpaligned_alloc(Alignment, Size);
	}
	else
	{
		NewPtr = rpaligned_alloc(uint32_t(Size >= 16 ? 16 : 8), Size);
	}

#	if ZEN_BUILD_DEBUG && ZEN_ENABLE_DEBUG_FILL
	if (Size && NewPtr != nullptr)
	{
		memset(NewPtr, DEBUG_FILL_NEW, rpmalloc_usable_size(NewPtr));
	}
#	endif

	return NewPtr;
}

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

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

	return Result;
}

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

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

	return Result;
}

void*
FMallocRpmalloc::TryRealloc(void* Ptr, size_t NewSize, uint32_t Alignment)
{
#	if ZEN_BUILD_DEBUG && ZEN_ENABLE_DEBUG_FILL
	size_t OldSize = 0;
	if (Ptr)
	{
		OldSize = rpmalloc_usable_size(Ptr);
		if (NewSize < OldSize)
		{
			memset((uint8_t*)Ptr + NewSize, DEBUG_FILL_FREED, OldSize - NewSize);
		}
	}
#	endif
	void* NewPtr = nullptr;

	if (NewSize == 0)
	{
		rpfree(Ptr);

		return nullptr;
	}

#	if ZEN_PLATFORM_MAC
	// macOS expects all allocations to be aligned to 16 bytes, so on Mac we always have to use mi_realloc_aligned
	Alignment = AlignArbitrary(Max((uint32_t)16, Alignment), (uint32_t)16);
	NewPtr	  = rpaligned_realloc(Ptr, Alignment, NewSize, /* OldSize */ 0, /* flags */ 0);
#	else
	if (Alignment != DEFAULT_ALIGNMENT)
	{
		Alignment = Max(NewSize >= 16 ? (uint32_t)16 : (uint32_t)8, Alignment);
		NewPtr	  = rpaligned_realloc(Ptr, Alignment, NewSize, /* OldSize */ 0, /* flags */ 0);
	}
	else
	{
		NewPtr = rprealloc(Ptr, NewSize);
	}
#	endif

#	if ZEN_BUILD_DEBUG && ZEN_ENABLE_DEBUG_FILL
	if (NewPtr && NewSize > OldSize)
	{
		memset((uint8_t*)NewPtr + OldSize, DEBUG_FILL_NEW, rpmalloc_usable_size(NewPtr) - OldSize);
	}
#	endif

	return NewPtr;
}

void
FMallocRpmalloc::Free(void* Ptr)
{
	if (!Ptr)
	{
		return;
	}

#	if ZEN_BUILD_DEBUG && ZEN_ENABLE_DEBUG_FILL
	memset(Ptr, DEBUG_FILL_FREED, rpmalloc_usable_size(Ptr));
#	endif

	rpfree(Ptr);
}

void*
FMallocRpmalloc::MallocZeroed(size_t Size, uint32_t Alignment)
{
	void* Result = TryMallocZeroed(Size, Alignment);

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

	return Result;
}
void*
FMallocRpmalloc::TryMallocZeroed(size_t Size, uint32_t Alignment)
{
	void* NewPtr = nullptr;

	if (Alignment != DEFAULT_ALIGNMENT)
	{
		Alignment = Max(uint32_t(Size >= 16 ? 16 : 8), Alignment);
		NewPtr	  = rpaligned_zalloc(Alignment, Size);
	}
	else
	{
		NewPtr = rpaligned_zalloc(uint32_t(Size >= 16 ? 16 : 8), Size);
	}

	return NewPtr;
}

bool
FMallocRpmalloc::GetAllocationSize(void* Original, size_t& SizeOut)
{
	// this is not the same as the allocation size - is that ok?
	SizeOut = rpmalloc_usable_size(Original);
	return true;
}
void
FMallocRpmalloc::Trim(bool bTrimThreadCaches)
{
	ZEN_UNUSED(bTrimThreadCaches);
}
}  // namespace zen
#endif