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,26 +15,72 @@ 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 PLATFORM_WINDOWS
-	if (Count != 0) Result = _aligned_malloc(Count, Alignment);
-#else
-	void* Ptr = SystemMalloc(Count + Alignment + sizeof(void*) + sizeof(size_t));
-	if (Ptr)
 	{
-		Result = Align(reinterpret_cast<uint8*>(Ptr) + sizeof(void*) + sizeof(size_t), Alignment);
+		Result = _aligned_malloc(Count, 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
+#	else
+	{
+		void* Ptr = SystemMalloc(Count + Alignment + sizeof(void*) + sizeof(size_t));
+
+		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
+#		if PLATFORM_WINDOWS
-		return _aligned_realloc(Ptr, Count, Alignment);
+		{
-#else
+			Result = _aligned_realloc(Ptr, Count, Alignment);
-		void* Result = Malloc(Count, Alignment);
+		}
-		size_t PtrSize = *reinterpret_cast<size_t*>(reinterpret_cast<uint8*>(Ptr) - sizeof(void*) - sizeof(size_t));
+#		else
-		Memcpy(Result, Ptr, Count < PtrSize ? Count : PtrSize);
+		{
-		Free(Ptr);
+			Result = Malloc(Count, Alignment);
-		return Result;
+
-#endif
+			if (Result != nullptr)
-	}
+			{
-	else if (Ptr == nullptr)
+				size_t PtrSize = *reinterpret_cast<size_t*>(reinterpret_cast<uint8*>(Ptr) - sizeof(void*) - sizeof(size_t));
-	{
+				Memcpy(Result, Ptr, Count < PtrSize ? Count : PtrSize);
-		return Malloc(Count, Alignment);
+				Free(Ptr);
+			}
+		}
+#		endif
 	}
 	else
 	{
-		Free(Ptr);
+		Result = Malloc(Count, Alignment);
-		return nullptr;
 	}
+
+	check(Result != nullptr);
+
+	return Result;
 }
 
 void Free(void* Ptr)
 {
-#if PLATFORM_WINDOWS
+#	if PLATFORM_WINDOWS
-	_aligned_free(Ptr);
+	{
-#else
+		_aligned_free(Ptr);
-	SystemFree(*reinterpret_cast<void**>(reinterpret_cast<uint8*>(Ptr) - sizeof(void*)));
+	}
-#endif
+#	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.