From 7cb93558053dedae78b68c6f003a1cea62ce92a4 Mon Sep 17 00:00:00 2001
From: Jamie Reece Wilson <me@reece.sx>
Date: Mon, 15 Apr 2024 17:51:22 +0100
Subject: [PATCH] [*] Harden RequestHeapOfRegion

---
 Include/Aurora/Memory/MemoryView.hpp |  7 +++++
 Source/Memory/AuHeapInternal.cpp     | 46 +++++++++++++++++-----------
 Source/Memory/AuHeapInternal.hpp     |  8 +++--
 3 files changed, 40 insertions(+), 21 deletions(-)

diff --git a/Include/Aurora/Memory/MemoryView.hpp b/Include/Aurora/Memory/MemoryView.hpp
index 097eed0f..0f1620a7 100644
--- a/Include/Aurora/Memory/MemoryView.hpp
+++ b/Include/Aurora/Memory/MemoryView.hpp
@@ -354,6 +354,13 @@ namespace Aurora::Memory
 
             AuResetMember(this->controlBlock.pPinner);
         }
+
+    public:
+        bool HasControlBlock() const
+        {
+            return this->controlBlock.pInUseCounter ||
+                this->controlBlock.pPinner;
+        }
     };
 
     using MemoryViewRead  = MemoryView<true>;
diff --git a/Source/Memory/AuHeapInternal.cpp b/Source/Memory/AuHeapInternal.cpp
index 7b87aa5c..13469bcc 100644
--- a/Source/Memory/AuHeapInternal.cpp
+++ b/Source/Memory/AuHeapInternal.cpp
@@ -22,13 +22,13 @@ typedef struct FragmentHeader
 namespace Aurora::Memory
 {
     InternalHeap::InternalHeap() :
-        heap_(nullptr),
-        count_(0)
+        pHeap_(nullptr),
+        uAllocCount_(0)
     { }
 
     InternalHeap::InternalHeap(const MemoryViewWrite &memory)
     {
-        SysAssert(this->Init(memory.length, memory.ptr));
+        SysAssert(this->Init(memory));
     }
 
     InternalHeap::InternalHeap(AuUInt uLength)
@@ -38,13 +38,13 @@ namespace Aurora::Memory
 
     InternalHeap::~InternalHeap()
     {
-        SysAssertDbgExp(this->count_ == 0);
+        SysAssertDbgExp(this->uAllocCount_ == 0);
 
         if (this->pBase_)
         {
-            if (this->heap_)
+            if (this->pHeap_)
             {
-                this->heap_ = nullptr;
+                this->pHeap_ = nullptr;
             }
 
             if (this->bOwnsMemory_)
@@ -60,6 +60,16 @@ namespace Aurora::Memory
         return reinterpret_cast<const FragmentHeader *>(ptr)[-1].size;
     }
 
+    bool InternalHeap::Init(const MemoryViewWrite &memory)
+    {
+        if (memory.HasControlBlock())
+        {
+            this->pBloat_ = AuMakeShared<MemoryViewWrite>(memory);
+        }
+
+        return this->Init(memory.length, memory.ptr);
+    }
+
     bool InternalHeap::Init(AuUInt uLength, void *ptr)
     {
         SysAssert(!this->pBase_, "heap already initialized");
@@ -87,7 +97,7 @@ namespace Aurora::Memory
             this->uLength_     = uLength;
         }
 
-        if (!(this->heap_ = o1heapInit(this->pBase_, uLength)))
+        if (!(this->pHeap_ = o1heapInit(this->pBase_, uLength)))
         {
             return false;
         }
@@ -130,15 +140,15 @@ namespace Aurora::Memory
 
     void *InternalHeap::_FAlloc(Types::size_t uLength)
     {
-        if (!this->heap_)
+        if (!this->pHeap_)
         {
             return nullptr;
         }
 
-        auto ret = o1heapAllocate(this->heap_, uLength);
+        auto ret = o1heapAllocate(this->pHeap_, uLength);
         if (ret)
         {
-            AuAtomicAdd(&this->count_, 1);
+            AuAtomicAdd(&this->uAllocCount_, 1u);
         }
         return ret;
     }
@@ -151,7 +161,7 @@ namespace Aurora::Memory
 
     void *InternalHeap::_ZAlloc(Types::size_t uLength)
     {
-        if (!this->heap_)
+        if (!this->pHeap_)
         {
             return nullptr;
         }
@@ -219,13 +229,13 @@ namespace Aurora::Memory
             return;
         }
 
-        o1heapFree(this->heap_, pBuffer);
+        o1heapFree(this->pHeap_, pBuffer);
         DecrementUsers();
     }
 
     void InternalHeap::DecrementUsers()
     {
-        if (AuAtomicSub(&this->count_, 1) == 0)
+        if (AuAtomicSub(&this->uAllocCount_, 1u) == 0)
         {
             TryRelease();
         }
@@ -238,7 +248,7 @@ namespace Aurora::Memory
             return;
         }
 
-        if (count_ == 0)
+        if (AuAtomicLoad(&uAllocCount_) == 0)
         {
             delete this;
         }
@@ -246,7 +256,7 @@ namespace Aurora::Memory
 
     void InternalHeap::RequestTermination()
     {
-        if (AuAtomicLoad(&this->count_))
+        if (AuAtomicLoad(&this->uAllocCount_))
         {
             SysPushErrorMemory("Heap life was less than its allocations, waiting for final free");
             SysPushErrorMemory("Reporting using mayday!");
@@ -265,7 +275,7 @@ namespace Aurora::Memory
 
     void InternalHeap::UpdateStats()
     {
-        auto pDiag = o1heapGetDiagnostics(this->heap_);
+        auto pDiag = o1heapGetDiagnostics(this->pHeap_);
 
         this->stats.uBytesLiveCounter = pDiag.allocated;
         this->stats.uBytesCapacity    = pDiag.capacity;
@@ -274,7 +284,7 @@ namespace Aurora::Memory
 
     void InternalHeap::WalkHeap(bool(*fCallback)(void *, void *), void *pSecondArg)
     {
-        o1heapTraverseHeap(this->heap_, fCallback, pSecondArg);
+        o1heapTraverseHeap(this->pHeap_, fCallback, pSecondArg);
     }
 
     AuSPtr<Heap> InternalHeap::GetSelfReference()
@@ -330,7 +340,7 @@ namespace Aurora::Memory
             return nullptr;
         }
 
-        if (!pHeap->Init(memory.length, memory.ptr))
+        if (!pHeap->Init(memory))
         {
             delete pHeap;
             return nullptr;
diff --git a/Source/Memory/AuHeapInternal.hpp b/Source/Memory/AuHeapInternal.hpp
index 88a96b30..266c3969 100644
--- a/Source/Memory/AuHeapInternal.hpp
+++ b/Source/Memory/AuHeapInternal.hpp
@@ -26,6 +26,7 @@ namespace Aurora::Memory
         InternalHeap(AuUInt uLength);
         virtual ~InternalHeap();
 
+        bool Init(const MemoryViewWrite &memory);
         bool Init(AuUInt uLength, void *ptr = nullptr);
 
         static AuUInt GetHeapSize(const void *ptr);
@@ -53,12 +54,13 @@ namespace Aurora::Memory
         void WalkHeap(bool(*fCallback)(void *, void *), void *pSecondArg) override;
 
     private:
-        O1HeapInstance *heap_ {};
-        int count_ {};
-        bool bIsDangling_ {};
+        O1HeapInstance *pHeap_ {};
+        std::shared_ptr<MemoryViewWrite> pBloat_;
     
     protected:
+        AuUInt32 uAllocCount_ {};
         bool bOwnsMemory_ {};
+        bool bIsDangling_ {};
 
     };
 }
\ No newline at end of file