From fe36ee5fb2b8f47aa926b3223c18d6cba1c6abf8 Mon Sep 17 00:00:00 2001
From: Jamie Reece Wilson <me@reece.sx>
Date: Tue, 17 Oct 2023 08:45:44 +0100
Subject: [PATCH] [+] AuMemory::SetLeakFinder

---
 Include/Aurora/Memory/Memory.hpp |   6 ++
 Source/Memory/Memory.cpp         | 107 ++++++++++++++++++++++++++++---
 2 files changed, 105 insertions(+), 8 deletions(-)

diff --git a/Include/Aurora/Memory/Memory.hpp b/Include/Aurora/Memory/Memory.hpp
index 9456e1c4..6c8ccc11 100644
--- a/Include/Aurora/Memory/Memory.hpp
+++ b/Include/Aurora/Memory/Memory.hpp
@@ -20,6 +20,12 @@
 
 namespace Aurora::Memory
 {
+    using LeakFinderAlloc_f = void(__cdecl *)(void *, AuUInt);
+    using LeakFinderFree_f = void(__cdecl *)(void *);
+
+    AUKN_SYM void SetLeakFinder(LeakFinderAlloc_f pAlloc,
+                                LeakFinderFree_f pFree);
+
     AUKN_SYM AU_ALLOC void *_ZAlloc(Types::size_t length);
     AUKN_SYM AU_ALLOC void *_ZAlloc(Types::size_t length, Types::size_t align);
     AUKN_SYM AU_ALLOC void *_FAlloc(Types::size_t length);
diff --git a/Source/Memory/Memory.cpp b/Source/Memory/Memory.cpp
index 32d9b61f..a5e99145 100644
--- a/Source/Memory/Memory.cpp
+++ b/Source/Memory/Memory.cpp
@@ -22,6 +22,9 @@ namespace Aurora::Memory
     AuUInt gBytesCounterPeak {};
 
     thread_local AuInt64 tlsLastOutOfMemory {};
+    
+    static LeakFinderAlloc_f gLeakFinderAlloc;
+    static LeakFinderFree_f gLeakFinderFree;
 
     static void AddBytesToCounter(AuUInt uBytes)
     {
@@ -32,6 +35,13 @@ namespace Aurora::Memory
     {
         AuAtomicSub(&gBytesCounterAllocated, uBytes);
     }
+    
+    AUKN_SYM void SetLeakFinder(LeakFinderAlloc_f pAlloc,
+                                LeakFinderFree_f pFree)
+    {
+        gLeakFinderAlloc = pAlloc;
+        gLeakFinderFree = pFree;
+    }
 
     AUKN_SYM AuUInt GetChunkSize(const void *head)
     {
@@ -73,6 +83,10 @@ namespace Aurora::Memory
         {                                                                                        \
             AddBytesToCounter(::mi_malloc_size(pRet));                                           \
         }                                                                                        \
+        if (gLeakFinderAlloc)                                                                    \
+        {                                                                                        \
+            gLeakFinderAlloc(pRet, ::mi_malloc_size(pRet));                                      \
+        }                                                                                        \
         return pRet;
 
     AUKN_SYM void *_ZAlloc(Types::size_t length)
@@ -110,9 +124,27 @@ namespace Aurora::Memory
         }
         else
         {
-            RemoveBytesFromCounter(::mi_malloc_size(buffer));
+            auto oldLen = ::mi_malloc_size(buffer);
+
             pRet = ::mi_rezalloc_aligned(buffer, length, align);
-            AddBytesToCounter(::mi_malloc_size(pRet));
+
+            if (pRet)
+            {
+                auto uNewSize = ::mi_malloc_size(pRet);
+
+                RemoveBytesFromCounter(oldLen);
+                AddBytesToCounter(uNewSize);
+                
+                if (gLeakFinderAlloc)
+                {
+                    if (gLeakFinderFree)
+                    {
+                        gLeakFinderFree(buffer);
+                    }
+
+                    gLeakFinderAlloc(pRet, uNewSize);
+                }
+            }
         }
 
         if (!pRet)
@@ -137,9 +169,27 @@ namespace Aurora::Memory
         }
         else
         {
-            RemoveBytesFromCounter(::mi_malloc_size(buffer));
+            auto oldLen = ::mi_malloc_size(buffer);
+
             pRet = ::mi_rezalloc(buffer, length);
-            AddBytesToCounter(::mi_malloc_size(pRet));
+
+            if (pRet)
+            {
+                auto uNewSize = ::mi_malloc_size(pRet);
+
+                RemoveBytesFromCounter(oldLen);
+                AddBytesToCounter(uNewSize);
+                
+                if (gLeakFinderAlloc)
+                {
+                    if (gLeakFinderFree)
+                    {
+                        gLeakFinderFree(buffer);
+                    }
+
+                    gLeakFinderAlloc(pRet, uNewSize);
+                }
+            }
         }
 
         if (!pRet)
@@ -164,9 +214,27 @@ namespace Aurora::Memory
         }
         else
         {
-            RemoveBytesFromCounter(::mi_malloc_size(buffer));
+            auto oldLen = ::mi_malloc_size(buffer);
+
             pRet = ::mi_realloc_aligned(buffer, length, align);
-            AddBytesToCounter(::mi_malloc_size(pRet));
+
+            if (pRet)
+            {
+                auto uNewSize = ::mi_malloc_size(pRet);
+
+                RemoveBytesFromCounter(oldLen);
+                AddBytesToCounter(uNewSize);
+                
+                if (gLeakFinderAlloc)
+                {
+                    if (gLeakFinderFree)
+                    {
+                        gLeakFinderFree(buffer);
+                    }
+
+                    gLeakFinderAlloc(pRet, uNewSize);
+                }
+            }
         }
 
         if (!pRet)
@@ -191,9 +259,27 @@ namespace Aurora::Memory
         }
         else
         {
-            RemoveBytesFromCounter(::mi_malloc_size(buffer));
+            auto oldLen = ::mi_malloc_size(buffer);
+
             pRet = ::mi_realloc(buffer, length);
-            AddBytesToCounter(::mi_malloc_size(pRet));
+
+            if (pRet)
+            {
+                auto uNewSize = ::mi_malloc_size(pRet);
+
+                RemoveBytesFromCounter(oldLen);
+                AddBytesToCounter(uNewSize);
+                
+                if (gLeakFinderAlloc)
+                {
+                    if (gLeakFinderFree)
+                    {
+                        gLeakFinderFree(buffer);
+                    }
+
+                    gLeakFinderAlloc(pRet, uNewSize);
+                }
+            }
         }
 
         if (!pRet)
@@ -223,6 +309,11 @@ namespace Aurora::Memory
             RemoveBytesFromCounter(::mi_malloc_size(pHead));
             ::mi_free(pHead);
         }
+
+        if (gLeakFinderFree)
+        {
+            gLeakFinderFree(pHead);
+        }
     }
 
     AUKN_SYM AuInt64 GetLastOutOfMemoryTimeNS()