From b58d352eaeca167ceef9d844479132863b165c0c Mon Sep 17 00:00:00 2001
From: Aiden <vabold820@gmail.com>
Date: Sat, 12 Oct 2024 03:40:57 -0400
Subject: [PATCH] Add heaps to scenes

---
 source/egg/core/Disposer.hh           |  6 +++++
 source/egg/core/ExpHeap.cc            | 19 -------------
 source/egg/core/ExpHeap.hh            |  6 -----
 source/egg/core/Heap.cc               | 30 ++++-----------------
 source/egg/core/Heap.hh               |  6 +----
 source/egg/core/Scene.cc              |  7 +++++
 source/egg/core/Scene.hh              |  2 ++
 source/egg/core/SceneManager.cc       | 39 ++++++++++++++++++++++-----
 source/egg/core/SceneManager.hh       |  7 +++++
 source/game/kart/KartObjectManager.cc |  5 ++++
 source/game/kart/KartObjectProxy.cc   |  4 +++
 source/game/kart/KartObjectProxy.hh   |  2 ++
 source/game/scene/GameScene.cc        |  1 +
 source/game/scene/RaceScene.cc        |  4 ++-
 source/game/scene/RootScene.cc        |  4 ++-
 source/host/System.cc                 | 28 ++++++++++++++++++-
 source/host/System.hh                 |  4 +++
 17 files changed, 110 insertions(+), 64 deletions(-)

diff --git a/source/egg/core/Disposer.hh b/source/egg/core/Disposer.hh
index 566afc4c..79e93d6a 100644
--- a/source/egg/core/Disposer.hh
+++ b/source/egg/core/Disposer.hh
@@ -11,12 +11,18 @@ class Heap;
 class Disposer {
     friend class Heap;
 
+public:
+    [[nodiscard]] static constexpr u16 getLinkOffset() {
+        return reinterpret_cast<uintptr_t>(&reinterpret_cast<Disposer *>(NULL)->m_link);
+    }
+
 protected:
     Disposer();
     virtual ~Disposer();
 
 private:
     Heap *m_heap;
+    Abstract::Memory::MEMLink m_link;
 };
 
 } // namespace EGG
diff --git a/source/egg/core/ExpHeap.cc b/source/egg/core/ExpHeap.cc
index f2204d4f..a5188099 100644
--- a/source/egg/core/ExpHeap.cc
+++ b/source/egg/core/ExpHeap.cc
@@ -125,23 +125,6 @@ const MEMiExpHeapHead *ExpHeap::dynamicCastHandleToExp() const {
     return reinterpret_cast<MEMiExpHeapHead *>(m_handle);
 }
 
-void ExpHeap::initRootHeap(void *startAddress, size_t size) {
-    MEMiHeapHead::OptFlag opt;
-    opt.setBit(MEMiHeapHead::eOptFlag::ZeroFillAlloc);
-
-#ifdef BUILD_DEBUG
-    opt.setBit(MEMiHeapHead::eOptFlag::DebugFillAlloc);
-#endif
-
-    s_rootHeap = create(startAddress, size, 2);
-    s_rootHeap->setName("EGGRoot");
-    s_rootHeap->becomeCurrentHeap();
-}
-
-ExpHeap *ExpHeap::getRootHeap() {
-    return s_rootHeap;
-}
-
 /// @addr{0x80226DD0}
 ExpHeap::GroupSizeRecord::GroupSizeRecord() {
     reset();
@@ -163,6 +146,4 @@ void ExpHeap::GroupSizeRecord::addSize(u16 groupID, size_t size) {
     m_entries[groupID] += size;
 }
 
-ExpHeap *ExpHeap::s_rootHeap = nullptr;
-
 } // namespace EGG
diff --git a/source/egg/core/ExpHeap.hh b/source/egg/core/ExpHeap.hh
index 2b772a93..ba14248d 100644
--- a/source/egg/core/ExpHeap.hh
+++ b/source/egg/core/ExpHeap.hh
@@ -47,14 +47,8 @@ public:
     [[nodiscard]] static ExpHeap *create(void *startAddress, size_t size, u16 opt);
     [[nodiscard]] static ExpHeap *create(size_t size, Heap *heap, u16 opt);
 
-    static void initRootHeap(void *startAddress, size_t size);
-
-    [[nodiscard]] static ExpHeap *getRootHeap();
-
 private:
     ExpHeap(Abstract::Memory::MEMiHeapHead *handle);
-
-    static ExpHeap *s_rootHeap;
 };
 
 } // namespace EGG
diff --git a/source/egg/core/Heap.cc b/source/egg/core/Heap.cc
index f9646d6f..17361a1a 100644
--- a/source/egg/core/Heap.cc
+++ b/source/egg/core/Heap.cc
@@ -5,11 +5,7 @@ using namespace Abstract::Memory;
 namespace EGG {
 
 /// @addr{0x802296E8}
-Heap::Heap(MEMiHeapHead *handle) : m_handle(handle), m_children() {
-    if (!s_isHeapInitialized) {
-        PANIC("Cannot create a heap before calling Heap::initalize");
-    }
-
+Heap::Heap(MEMiHeapHead *handle) : m_handle(handle), m_children(Disposer::getLinkOffset()) {
     m_block = nullptr;
     m_parentHeap = nullptr;
     m_name = "NoName";
@@ -25,12 +21,12 @@ Heap::~Heap() {
 
 /// @addr{0x80229C5C}
 void Heap::dispose() {
-    for (auto *&node : m_children) {
+    Disposer *node = nullptr;
+    while ((node = reinterpret_cast<Disposer *>(m_children.getFirst()))) {
         node->~Disposer();
-        node = nullptr;
     }
 
-    m_children.clear();
+    ASSERT(!m_children.m_headObject && !m_children.m_tailObject);
 }
 
 void Heap::disableAllocation() {
@@ -46,7 +42,7 @@ bool Heap::tstDisableAllocation() const {
 }
 
 void Heap::appendDisposer(Disposer *disposer) {
-    m_children.push_back(disposer);
+    m_children.append(disposer);
 }
 
 void Heap::removeDisposer(Disposer *disposer) {
@@ -95,16 +91,6 @@ void Heap::setParentHeap(Heap *heap) {
     m_parentHeap = heap;
 }
 
-/// @addr{0x802296A8}
-void Heap::initialize() {
-    constexpr size_t MEMORY_SPACE_SIZE = 0x1000000;
-
-    s_isHeapInitialized = true;
-
-    s_memorySpace = malloc(MEMORY_SPACE_SIZE);
-    ExpHeap::initRootHeap(s_memorySpace, MEMORY_SPACE_SIZE);
-}
-
 /// @addr{0x80229814}
 void *Heap::alloc(size_t size, int align, Heap *pHeap) {
     Heap *currentHeap = s_currentHeap;
@@ -195,10 +181,6 @@ Heap *Heap::getCurrentHeap() {
     return s_currentHeap;
 }
 
-void *Heap::getMemorySpace() {
-    return s_memorySpace;
-}
-
 } // namespace EGG
 
 /// @addr{0x80229DCC}
@@ -252,6 +234,4 @@ void operator delete[](void *block, size_t /* size */) noexcept {
 MEMList EGG::Heap::s_heapList = MEMList(EGG::Heap::getOffset()); ///< @addr{0x80384320}
 
 EGG::Heap *EGG::Heap::s_currentHeap = nullptr;     ///< @addr{0x80386EA0}
-bool EGG::Heap::s_isHeapInitialized = false;       ///< @addr{0x80386EA4}
 EGG::Heap *EGG::Heap::s_allocatableHeap = nullptr; ///< @addr{0x80386EA8}
-void *EGG::Heap::s_memorySpace = nullptr;
diff --git a/source/egg/core/Heap.hh b/source/egg/core/Heap.hh
index 01d64750..6b1d6ceb 100644
--- a/source/egg/core/Heap.hh
+++ b/source/egg/core/Heap.hh
@@ -63,7 +63,6 @@ public:
 
     [[nodiscard]] static ExpHeap *dynamicCastToExp(Heap *heap);
     [[nodiscard]] static Heap *getCurrentHeap();
-    [[nodiscard]] static void *getMemorySpace();
 
     [[nodiscard]] static constexpr uintptr_t getOffset() {
         // offsetof doesn't work, so instead of hardcoding an offset, we derive it ourselves
@@ -81,16 +80,13 @@ protected:
     Heap *m_parentHeap;
     Flags m_flags;
     Abstract::Memory::MEMLink m_link;
-    std::list<Disposer *> m_children;
+    Abstract::Memory::MEMList m_children;
     const char *m_name;
 
     static Abstract::Memory::MEMList s_heapList;
 
     static Heap *s_currentHeap;
-    static bool s_isHeapInitialized;
     static Heap *s_allocatableHeap;
-    static ExpHeap *s_rootHeap;
-    static void *s_memorySpace;
 };
 
 } // namespace EGG
diff --git a/source/egg/core/Scene.cc b/source/egg/core/Scene.cc
index a3041101..43ae400b 100644
--- a/source/egg/core/Scene.cc
+++ b/source/egg/core/Scene.cc
@@ -1,9 +1,12 @@
 #include "Scene.hh"
 
+#include "egg/core/SceneManager.hh"
+
 namespace EGG {
 
 /// @addr{0x8023AD10}
 Scene::Scene() {
+    m_heap = SceneManager::heapForCreateScene();
     m_parent = nullptr;
     m_child = nullptr;
     m_id = -1;
@@ -29,6 +32,10 @@ void Scene::setSceneMgr(SceneManager *sceneMgr) {
     m_sceneMgr = sceneMgr;
 }
 
+Heap *Scene::heap() const {
+    return m_heap;
+}
+
 Scene *Scene::parent() const {
     return m_parent;
 }
diff --git a/source/egg/core/Scene.hh b/source/egg/core/Scene.hh
index d1a32088..7dd3cbeb 100644
--- a/source/egg/core/Scene.hh
+++ b/source/egg/core/Scene.hh
@@ -41,6 +41,7 @@ public:
     /// @endSetters
 
     /// @beginGetters
+    [[nodiscard]] Heap *heap() const;
     [[nodiscard]] Scene *parent() const;
     [[nodiscard]] Scene *child() const;
     [[nodiscard]] int id() const;
@@ -52,6 +53,7 @@ protected:
         Members
      *-----------*/
 
+    Heap *m_heap;
     Scene *m_parent;
     Scene *m_child;
     int m_id;
diff --git a/source/egg/core/SceneManager.cc b/source/egg/core/SceneManager.cc
index c362daaa..b2f0333a 100644
--- a/source/egg/core/SceneManager.cc
+++ b/source/egg/core/SceneManager.cc
@@ -1,5 +1,9 @@
 #include "SceneManager.hh"
 
+#include "egg/core/ExpHeap.hh"
+
+#include <host/System.hh>
+
 namespace EGG {
 
 /*------------*
@@ -84,6 +88,24 @@ void SceneManager::createChildScene(int id, Scene *parent) {
 
 /// @addr{0x8023B0E4}
 void SceneManager::createScene(int id, Scene *parent) {
+    Heap *parentHeap = parent ? parent->heap() : Host::KSystem::Instance().rootHeap();
+
+    // We need to preserve the locked status to reinstate it later
+    bool locked = parentHeap->tstDisableAllocation();
+    if (locked) {
+        parentHeap->enableAllocation();
+    }
+
+    ExpHeap *newHeap = ExpHeap::create(-1, parentHeap, sHeapOptionFlg);
+    sHeapForCreateScene = newHeap;
+
+    if (locked) {
+        parentHeap->disableAllocation();
+    }
+
+    newHeap->becomeCurrentHeap();
+    newHeap->setName("DefaultSceneHeap");
+
     Scene *newScene = m_creator->create(id);
 
     if (parent) {
@@ -101,19 +123,21 @@ void SceneManager::createScene(int id, Scene *parent) {
 void SceneManager::destroyScene(Scene *scene) {
     scene->exit();
     if (scene->child()) {
-        destroyScene(scene);
+        destroyScene(scene->child());
     }
 
     Scene *parent = scene->parent();
-    delete m_currentScene;
+    m_creator->destroy(scene->id());
     m_currentScene = nullptr;
 
-    if (!parent) {
-        return;
+    if (parent) {
+        parent->setChild(nullptr);
+        m_currentScene = parent;
     }
 
-    parent->setChild(nullptr);
-    m_currentScene = parent;
+    scene->heap()->destroy();
+    Heap *parentHeap = parent ? parent->heap() : Host::KSystem::Instance().rootHeap();
+    parentHeap->becomeCurrentHeap();
 }
 
 /// @addr{0x8023AF84}
@@ -257,4 +281,7 @@ void SceneManager::resetAfterFadeType() {
     m_fadeType = FadeType::Idle;
 }
 
+Heap *SceneManager::sHeapForCreateScene = nullptr;
+u16 SceneManager::sHeapOptionFlg = 2;
+
 } // namespace EGG
diff --git a/source/egg/core/SceneManager.hh b/source/egg/core/SceneManager.hh
index f7916694..8a95b576 100644
--- a/source/egg/core/SceneManager.hh
+++ b/source/egg/core/SceneManager.hh
@@ -73,6 +73,10 @@ public:
         return m_currentSceneId;
     }
 
+    [[nodiscard]] static Heap *heapForCreateScene() {
+        return sHeapForCreateScene;
+    }
+
     /*----------*
         Setters
      *----------*/
@@ -97,6 +101,9 @@ private:
     int m_nextSceneId;
     int m_currentSceneId;
     int m_prevSceneId;
+
+    static Heap *sHeapForCreateScene;
+    static u16 sHeapOptionFlg;
 };
 
 } // namespace EGG
diff --git a/source/game/kart/KartObjectManager.cc b/source/game/kart/KartObjectManager.cc
index a9c21756..fd8da83e 100644
--- a/source/game/kart/KartObjectManager.cc
+++ b/source/game/kart/KartObjectManager.cc
@@ -81,6 +81,11 @@ KartObjectManager::~KartObjectManager() {
     }
 
     delete[] m_objects;
+
+    // If the proxy list is not cleared when we're done with the KartObjectManager, the list's
+    // destructor calls delete on all of the links remaining in the list. Since the heaps are
+    // gone by that point, this results in a segmentation fault. So, we clear the links here.
+    KartObjectProxy::proxyList().clear();
 }
 
 KartObjectManager *KartObjectManager::s_instance = nullptr; ///< @addr{0x809C18F8}
diff --git a/source/game/kart/KartObjectProxy.cc b/source/game/kart/KartObjectProxy.cc
index 7a5b3822..b5ec84a5 100644
--- a/source/game/kart/KartObjectProxy.cc
+++ b/source/game/kart/KartObjectProxy.cc
@@ -338,6 +338,10 @@ std::pair<EGG::Vector3f, EGG::Vector3f> KartObjectProxy::getCannonPosRot() {
     return std::pair(cannonPos + local60 * temp0, cannonRot);
 }
 
+std::list<KartObjectProxy *> &KartObjectProxy::proxyList() {
+    return s_proxyList;
+}
+
 /// @addr{0x805901D0}
 void KartObjectProxy::apply(size_t idx) {
     m_accessor = KartObjectManager::Instance()->object(idx)->accessor();
diff --git a/source/game/kart/KartObjectProxy.hh b/source/game/kart/KartObjectProxy.hh
index 994c7e45..49a3724c 100644
--- a/source/game/kart/KartObjectProxy.hh
+++ b/source/game/kart/KartObjectProxy.hh
@@ -123,6 +123,8 @@ public:
     [[nodiscard]] u16 tireCount() const;
     [[nodiscard]] bool hasFloorCollision(const WheelPhysics *wheelPhysics) const;
     [[nodiscard]] std::pair<EGG::Vector3f, EGG::Vector3f> getCannonPosRot();
+
+    [[nodiscard]] static std::list<KartObjectProxy *> &proxyList();
     /// @endGetters
 
 protected:
diff --git a/source/game/scene/GameScene.cc b/source/game/scene/GameScene.cc
index d9bd48af..ee6b8516 100644
--- a/source/game/scene/GameScene.cc
+++ b/source/game/scene/GameScene.cc
@@ -9,6 +9,7 @@ namespace Scene {
 
 /// @addr{0x8051A1E0}
 GameScene::GameScene() {
+    m_heap->setName("DefaultGameSceneHeap");
     m_nextSceneId = -1;
 }
 
diff --git a/source/game/scene/RaceScene.cc b/source/game/scene/RaceScene.cc
index 365be5be..7a7da8bc 100644
--- a/source/game/scene/RaceScene.cc
+++ b/source/game/scene/RaceScene.cc
@@ -14,7 +14,9 @@
 namespace Scene {
 
 /// @addr{0x80553B88}
-RaceScene::RaceScene() = default;
+RaceScene::RaceScene() {
+    m_heap->setName("RaceSceneHeap");
+}
 
 /// @addr{0x80553BD4}
 RaceScene::~RaceScene() = default;
diff --git a/source/game/scene/RootScene.cc b/source/game/scene/RootScene.cc
index cb748ae2..5cd7994e 100644
--- a/source/game/scene/RootScene.cc
+++ b/source/game/scene/RootScene.cc
@@ -10,7 +10,9 @@
 namespace Scene {
 
 /// @addr{0x80542878}
-RootScene::RootScene() = default;
+RootScene::RootScene() {
+    m_heap->setName("RootSceneHeap");
+}
 
 /// @addr{0x805429A8}
 RootScene::~RootScene() = default;
diff --git a/source/host/System.cc b/source/host/System.cc
index f30f0278..54deb6d6 100644
--- a/source/host/System.cc
+++ b/source/host/System.cc
@@ -2,6 +2,8 @@
 
 #include "host/SceneCreatorDynamic.hh"
 
+#include <egg/core/ExpHeap.hh>
+
 #include <abstract/File.hh>
 
 #include <cstring>
@@ -11,7 +13,7 @@ namespace Host {
 /// @brief The main entry point for the program.
 /// @addr{0x80008EF0}
 int KSystem::main(int argc, char **argv) {
-    EGG::Heap::initialize();
+    initMemory();
 
     if (argc < 2) {
         PANIC("Expected file argument");
@@ -103,6 +105,26 @@ void KSystem::init() {
     delete[] m_suiteData.data();
 }
 
+/// @addr{0x80242504}
+/// @brief Creates a memory space and a heap from that memory space.
+/// @warning This must be called first, or any `new` calls will have a nullptr heap!
+/// @details In the base game, the initial memory space is provided by the OS arena.
+/// For local implementations, we request memory from the OS via malloc, as it's not overwritten.
+void KSystem::initMemory() {
+    constexpr size_t MEMORY_SPACE_SIZE = 0x1000000;
+    Abstract::Memory::MEMiHeapHead::OptFlag opt;
+    opt.setBit(Abstract::Memory::MEMiHeapHead::eOptFlag::ZeroFillAlloc);
+
+#ifdef BUILD_DEBUG
+    opt.setBit(Abstract::Memory::MEMiHeapHead::eOptFlag::DebugFillAlloc);
+#endif
+
+    m_memorySpace = malloc(MEMORY_SPACE_SIZE);
+    m_rootHeap = EGG::ExpHeap::create(m_memorySpace, MEMORY_SPACE_SIZE, opt);
+    m_rootHeap->setName("EGGRoot");
+    m_rootHeap->becomeCurrentHeap();
+}
+
 /// @brief The main loop of the program.
 /// @return Whether or not all test cases have passed.
 /// @addr{0x8000951C}
@@ -115,6 +137,10 @@ bool KSystem::run() {
     return m_testDirector->sync();
 }
 
+EGG::Heap *KSystem::rootHeap() const {
+    return m_rootHeap;
+}
+
 const Test::TestDirector *KSystem::testDirector() const {
     return m_testDirector;
 }
diff --git a/source/host/System.hh b/source/host/System.hh
index a792a826..6e4615b2 100644
--- a/source/host/System.hh
+++ b/source/host/System.hh
@@ -20,8 +20,10 @@ public:
     [[nodiscard]] Option option(char *arg);
     void handleOption(Option opt, int argc, char **argv, int &i);
     void init();
+    void initMemory();
     [[nodiscard]] bool run();
 
+    [[nodiscard]] EGG::Heap *rootHeap() const;
     [[nodiscard]] const Test::TestDirector *testDirector() const;
 
     [[nodiscard]] static KSystem &Instance();
@@ -32,6 +34,8 @@ private:
     KSystem(KSystem &&) = delete;
     ~KSystem();
 
+    void *m_memorySpace;
+    EGG::Heap *m_rootHeap;
     std::span<u8> m_suiteData;
     EGG::SceneManager *m_sceneMgr;
     Test::TestDirector *m_testDirector;