From 977ff07c85a0dd4db08f4e8b60a991b715f5a2af Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?L=C3=A9o=20Lam?= <leo@leolam.fr>
Date: Fri, 4 Sep 2020 22:38:06 +0200
Subject: [PATCH] ksys: Implement more heap utilities

---
 data/uking_functions.csv          | 10 ++++----
 src/KingSystem/Utils/HeapUtil.cpp | 42 +++++++++++++++++++++++++++++++
 src/KingSystem/Utils/HeapUtil.h   |  6 +++++
 3 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/data/uking_functions.csv b/data/uking_functions.csv
index 507d58af..042f27bf 100644
--- a/data/uking_functions.csv
+++ b/data/uking_functions.csv
@@ -90336,11 +90336,11 @@
 0x00000071011f321c,sub_71011F321C,24,
 0x00000071011f3234,sub_71011F3234,24,
 0x00000071011f324c,sinitStageSelectMenu,172,
-0x00000071011f32f8,getHeapOrCurrentHeap,48,
-0x00000071011f3328,adjustHeap,44,
-0x00000071011f3354,getCurrentHeap,16,
-0x00000071011f3364,getDebugHeap,32,
-0x00000071011f3384,getDebugHeap2,48,
+0x00000071011f32f8,getHeapOrCurrentHeap,48,_ZN4ksys4util20getHeapOrCurrentHeapEPN4sead4HeapE
+0x00000071011f3328,adjustHeap,44,_ZN4ksys4util23adjustHeapOrCurrentHeapEPN4sead4HeapE
+0x00000071011f3354,getCurrentHeap,16,_ZN4ksys4util14getCurrentHeapEv
+0x00000071011f3364,getDebugHeap,32,_ZN4ksys4util12getDebugHeapEv
+0x00000071011f3384,getDebugHeap2,48,_ZN4ksys4util13getDebugHeap2Ev
 0x00000071011f33b4,sub_71011F33B4,32,
 0x00000071011f33d4,sub_71011F33D4,88,
 0x00000071011f342c,sub_71011F342C,72,
diff --git a/src/KingSystem/Utils/HeapUtil.cpp b/src/KingSystem/Utils/HeapUtil.cpp
index fb06c48f..30118e44 100644
--- a/src/KingSystem/Utils/HeapUtil.cpp
+++ b/src/KingSystem/Utils/HeapUtil.cpp
@@ -4,6 +4,26 @@
 
 namespace ksys::util {
 
+// Name, layout and purpose unknown; this is stubbed in release builds
+class DebugHeapHolder {
+    SEAD_SINGLETON_DISPOSER(DebugHeapHolder)
+public:
+    virtual ~DebugHeapHolder() = default;
+    virtual sead::Heap* getHeap() { return nullptr; }
+};
+
+SEAD_SINGLETON_DISPOSER_IMPL(DebugHeapHolder)
+
+// Name, layout and purpose unknown; this is stubbed in release builds
+class DebugHeapHolder2 {
+    SEAD_SINGLETON_DISPOSER(DebugHeapHolder2)
+public:
+    virtual ~DebugHeapHolder2() = default;
+    virtual sead::Heap* getHeap() { return nullptr; }
+};
+
+SEAD_SINGLETON_DISPOSER_IMPL(DebugHeapHolder2)
+
 sead::Heap* getHeapOrCurrentHeap(sead::Heap* heap) {
     if (heap)
         return heap;
@@ -11,7 +31,29 @@ sead::Heap* getHeapOrCurrentHeap(sead::Heap* heap) {
     if (!sead::ThreadMgr::instance())
         return nullptr;
 
+    return getCurrentHeap();
+}
+
+void adjustHeapOrCurrentHeap(sead::Heap* heap) {
+    if (!heap)
+        heap = getCurrentHeap();
+    heap->adjust();
+}
+
+sead::Heap* getCurrentHeap() {
     return sead::HeapMgr::instance()->getCurrentHeap();
 }
 
+sead::Heap* getDebugHeap() {
+    if (!DebugHeapHolder::instance())
+        return nullptr;
+    return DebugHeapHolder::instance()->getHeap();
+}
+
+sead::Heap* getDebugHeap2() {
+    if (DebugHeapHolder2::instance())
+        return DebugHeapHolder2::instance()->getHeap();
+    return getDebugHeap();
+}
+
 }  // namespace ksys::util
diff --git a/src/KingSystem/Utils/HeapUtil.h b/src/KingSystem/Utils/HeapUtil.h
index 0e6be692..599791ef 100644
--- a/src/KingSystem/Utils/HeapUtil.h
+++ b/src/KingSystem/Utils/HeapUtil.h
@@ -11,4 +11,10 @@ namespace ksys::util {
 /// Returns the specified heap if it is not null. Otherwise, this returns the current heap.
 sead::Heap* getHeapOrCurrentHeap(sead::Heap* heap);
 
+void adjustHeapOrCurrentHeap(sead::Heap* heap);
+
+sead::Heap* getCurrentHeap();
+sead::Heap* getDebugHeap();
+sead::Heap* getDebugHeap2();
+
 }  // namespace ksys::util