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

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

#if ZEN_MIMALLOC_ENABLED

#	include <mimalloc.h>

/** 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 {

// Dramatically reduce memory zeroing and page faults during alloc intense workloads
// by keeping freed pages for a little while instead of releasing them
// right away to the OS, effectively acting like a scratch buffer
// until pages are both freed and inactive for the delay specified
// in milliseconds.
int32_t GMiMallocMemoryResetDelay = 10000;

FMallocMimalloc::FMallocMimalloc()
{
	mi_option_set(mi_option_reset_delay, GMiMallocMemoryResetDelay);
}

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

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

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

	return NewPtr;
}

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

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

	return Result;
}

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

	if (NewSize == 0)
	{
		mi_free(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	  = mi_realloc_aligned(Ptr, NewSize, Alignment);
#	else
	if (Alignment != DEFAULT_ALIGNMENT)
	{
		Alignment = Max(NewSize >= 16 ? (uint32_t)16 : (uint32_t)8, Alignment);
		NewPtr	  = mi_realloc_aligned(Ptr, NewSize, Alignment);
	}
	else
	{
		NewPtr = mi_realloc(Ptr, NewSize);
	}
#	endif

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

	return NewPtr;
}

void*
FMallocMimalloc::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
FMallocMimalloc::Free(void* Ptr)
{
	if (!Ptr)
	{
		return;
	}

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

	mi_free(Ptr);
}

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

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

	return Result;
}

void*
FMallocMimalloc::TryMallocZeroed(size_t Size, uint32_t Alignment)
{
	void* NewPtr = nullptr;

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

	return NewPtr;
}

bool
FMallocMimalloc::GetAllocationSize(void* Original, size_t& SizeOut)
{
	SizeOut = mi_malloc_size(Original);
	return true;
}

void
FMallocMimalloc::Trim(bool bTrimThreadCaches)
{
	mi_collect(bTrimThreadCaches);
}

#	undef DEBUG_FILL_FREED
#	undef DEBUG_FILL_NEW

}  // namespace zen

#endif	// MIMALLOC_ENABLED