feat(memory): add memory leak check assertion

Redcraft.Utility/Source/Private/Memory/Memory.cpp

@@ -1,6 +1,8 @@
 #include "Memory/Memory.h"
 
 #include "Memory/Alignment.h"
+#include "Templates/Atomic.h"
+#include "Templates/ScopeHelper.h"
 #include "Miscellaneous/AssertionMacros.h"
 
 #if PLATFORM_WINDOWS
@@ -13,27 +15,73 @@ NAMESPACE_MODULE_BEGIN(Utility)
 
 NAMESPACE_BEGIN(Memory)
 
+#if DO_CHECK
+
+class FMemoryLeakChecker
+{
+private:
+
+	TAtomic<size_t> MemoryAllocationCount;
+
+public:
+
+	FORCEINLINE constexpr FMemoryLeakChecker()
+		: MemoryAllocationCount(0)
+	{ }
+
+	FORCEINLINE ~FMemoryLeakChecker()
+	{
+		checkf(MemoryAllocationCount.Load() == 0, TEXT("There is unfree memory. Please check for memory leaks."));
+	}
+
+	FORCEINLINE void AddMemoryAllocationCount()
+	{
+		MemoryAllocationCount.FetchAdd(1, EMemoryOrder::Relaxed);
+	}
+
+	FORCEINLINE void ReleaseMemoryAllocationCount()
+	{
+		MemoryAllocationCount.FetchSub(1, EMemoryOrder::Relaxed);
+	}
+
+};
+
+FMemoryLeakChecker MemoryLeakChecker;
+
+#endif
+
 void* Malloc(size_t Count, size_t Alignment)
 {
 	checkf(IsValidAlignment(Alignment), TEXT("The alignment value must be an integer power of 2."));
 
+	Count = Count != 0 ? Count : 1; // Treat zero-byte allocation as one-byte allocation.
+
 	const size_t MinimumAlignment = Count >= 16 ? 16 : 8;
 	Alignment = MinimumAlignment > Alignment ? MinimumAlignment : Alignment;
 
 	void* Result = nullptr;
 
 #	if PLATFORM_WINDOWS
-	if (Count != 0) Result = _aligned_malloc(Count, Alignment);
+	{
+		Result = _aligned_malloc(Count, Alignment);
+	}
 #	else
+	{
 		void* Ptr = SystemMalloc(Count + Alignment + sizeof(void*) + sizeof(size_t));
-	if (Ptr)
+
+		if (Ptr != nullptr)
 		{
 			Result = Align(reinterpret_cast<uint8*>(Ptr) + sizeof(void*) + sizeof(size_t), Alignment);
 			*reinterpret_cast<void**>(reinterpret_cast<uint8*>(Result) - sizeof(void*)) = Ptr;
 			*reinterpret_cast<size_t*>(reinterpret_cast<uint8*>(Result) - sizeof(void*) - sizeof(size_t)) = Count;
 		}
+	}
 #	endif
 
+	check(Result != nullptr);
+
+	check_code({ MemoryLeakChecker.AddMemoryAllocationCount(); });
+
 	return Result;
 }
 
@@ -41,39 +89,55 @@ void* Realloc(void* Ptr, size_t Count, size_t Alignment)
 {
 	checkf(IsValidAlignment(Alignment), TEXT("The alignment value must be an integer power of 2."));
 
+	Count = Count != 0 ? Count : 1; // Treat zero-byte allocation as one-byte allocation.
+
 	const size_t MinimumAlignment = Count >= 16 ? 16 : 8;
 	Alignment = MinimumAlignment > Alignment ? MinimumAlignment : Alignment;
 
-	if (Ptr && Count)
+	void* Result = nullptr;
+
+	if (Ptr != nullptr)
 	{
 #		if PLATFORM_WINDOWS
-		return _aligned_realloc(Ptr, Count, Alignment);
+		{
+			Result = _aligned_realloc(Ptr, Count, Alignment);
+		}
 #		else
-		void* Result = Malloc(Count, Alignment);
+		{
+			Result = Malloc(Count, Alignment);
+
+			if (Result != nullptr)
+			{
 				size_t PtrSize = *reinterpret_cast<size_t*>(reinterpret_cast<uint8*>(Ptr) - sizeof(void*) - sizeof(size_t));
 				Memcpy(Result, Ptr, Count < PtrSize ? Count : PtrSize);
 				Free(Ptr);
-		return Result;
-#endif
 			}
-	else if (Ptr == nullptr)
+		}
-	{
+#		endif
-		return Malloc(Count, Alignment);
 	}
 	else
 	{
-		Free(Ptr);
+		Result = Malloc(Count, Alignment);
-		return nullptr;
 	}
+
+	check(Result != nullptr);
+
+	return Result;
 }
 
 void Free(void* Ptr)
 {
 #	if PLATFORM_WINDOWS
+	{
 		_aligned_free(Ptr);
+	}
 #	else
+	{
 		SystemFree(*reinterpret_cast<void**>(reinterpret_cast<uint8*>(Ptr) - sizeof(void*)));
+	}
 #	endif
+
+	check_code({ if (Ptr != nullptr) MemoryLeakChecker.ReleaseMemoryAllocationCount(); });
 }
 
 size_t QuantizeSize(size_t Count, size_t Alignment)

Redcraft.Utility/Source/Private/Testing/MemoryTesting.cpp

@@ -148,6 +148,7 @@ void TestMemoryMalloc()
 	always_check(PtrC->A == 0x01234567);
 	delete [] PtrC;
 
+	Memory::Free(Memory::Realloc(Memory::Malloc(0), 0));
 }
 
 NAMESPACE_UNNAMED_BEGIN

Redcraft.Utility/Source/Public/Memory/Memory.h

@@ -255,6 +255,7 @@ NODISCARD REDCRAFTUTILITY_API void* Malloc(size_t Count, size_t Alignment = Defa
 
 /**
  * Reallocates the given area of memory. It must be previously allocated by Malloc() or Realloc().
+ * If 'Ptr' is a nullptr, effectively the same as calling Malloc().
  *
  * @param  Ptr       - The pointer to the memory area to be reallocated.
  * @param  Count     - The number of bytes to allocate.