dHeap(Allocator) work

5 files changed, 178 insertions, 44 deletions

diff --git a/config/SOUE01/splits.txt b/config/SOUE01/splits.txt
index f84fec00..5e1ad335 100644
--- a/config/SOUE01/splits.txt
+++ b/config/SOUE01/splits.txt
@@ -72,7 +72,9 @@ d/d_heap.cpp:
 	.text       start:0x80054B00 end:0x80054F30
 	.ctors      start:0x804DB680 end:0x804DB684
 	.rodata     start:0x804DE008 end:0x804DE188
-	.sbss       start:0x805751A8 end:0x805751D0
+	.data       start:0x8050D440 end:0x8050D458
+	.sbss       start:0x805751A8 end:0x805751CC
+	.bss        start:0x80597740 end:0x80597758
 
 d/d_main.cpp:
 	.text       start:0x80054F30 end:0x80055170
diff --git a/config/SOUE01/symbols.txt b/config/SOUE01/symbols.txt
index 195dcdc3..b345728e 100644
--- a/config/SOUE01/symbols.txt
+++ b/config/SOUE01/symbols.txt
@@ -2212,8 +2212,8 @@ fn_80054550 = .text:0x80054550; // type:function size:0x31C
 fn_80054870 = .text:0x80054870; // type:function size:0x25C
 fn_80054AD0 = .text:0x80054AD0; // type:function size:0x10
 fn_80054AE0 = .text:0x80054AE0; // type:function size:0x14
-fn_80054B00 = .text:0x80054B00; // type:function size:0x40
-fn_80054B40 = .text:0x80054B40; // type:function size:0x20
+__dt__14dHeapAllocatorFv = .text:0x80054B00; // type:function size:0x40
+onAlloc__14dHeapAllocatorFPQ23EGG12HeapAllocArg = .text:0x80054B40; // type:function size:0x20
 init__5dHeapFPCcUlPQ23EGG4Heap = .text:0x80054B60; // type:function size:0x64
 createWork1Heap__5dHeapFUlPQ23EGG4Heap = .text:0x80054BD0; // type:function size:0x18
 createWork2Heap__5dHeapFUlPQ23EGG4Heap = .text:0x80054BF0; // type:function size:0x18
@@ -2224,15 +2224,15 @@ createLayoutEx2Heap__5dHeapFUlPQ23EGG4Heap = .text:0x80054C70; // type:function
 createLayoutResHeap__5dHeapFUlPQ23EGG4Heap = .text:0x80054C90; // type:function size:0x18
 createFontHeap__5dHeapFUlPQ23EGG4Heap = .text:0x80054CB0; // type:function size:0x18
 createHBMHeap__5dHeapFUlPQ23EGG4Heap = .text:0x80054CD0; // type:function size:0x18
-fn_80054CF0 = .text:0x80054CF0; // type:function size:0x88
-fn_80054D80 = .text:0x80054D80; // type:function size:0x98
-fn_80054E20 = .text:0x80054E20; // type:function size:0x88
-fn_80054EB0 = .text:0x80054EB0; // type:function size:0x4
-fn_80054EC0 = .text:0x80054EC0; // type:function size:0xC
-fn_80054ED0 = .text:0x80054ED0; // type:function size:0xC
+initCallbacks__14dHeapAllocatorFv = .text:0x80054CF0; // type:function size:0x88
+allocCallback__18dHeapAllocatorBaseFPQ23EGG12HeapAllocArg = .text:0x80054D80; // type:function size:0x98
+freeCallback__18dHeapAllocatorBaseFPQ23EGG11HeapFreeArg = .text:0x80054E20; // type:function size:0x88
+onFree__18dHeapAllocatorBaseFPQ23EGG11HeapFreeArg = .text:0x80054EB0; // type:function size:0x4
+__nw__FUl = .text:0x80054EC0; // type:function size:0xC
+__nwa__FUl = .text:0x80054ED0; // type:function size:0xC
 __dl__FPv = .text:0x80054EE0; // type:function size:0x8
-fn_80054EF0 = .text:0x80054EF0; // type:function size:0x8
-fn_80054F00 = .text:0x80054F00; // type:function size:0x30
+__dla__FPv = .text:0x80054EF0; // type:function size:0x8
+__sinit_\d_heap_cpp = .text:0x80054F00; // type:function size:0x30
 fn_80054F30 = .text:0x80054F30; // type:function size:0x4
 fn_80054F40 = .text:0x80054F40; // type:function size:0x4
 fn_80054F50 = .text:0x80054F50; // type:function size:0x38
@@ -26205,11 +26205,11 @@ dumpAll__Q23EGG4HeapFv = .text:0x804958B0; // type:function size:0xEC
 becomeCurrentHeap__Q23EGG4HeapFv = .text:0x804959A0; // type:function size:0x58
 _becomeCurrentHeapWithoutLock__Q23EGG4HeapFv = .text:0x80495A00; // type:function size:0x3C
 initAllocator__Q23EGG4HeapFPQ23EGG9Allocatorl = .text:0x80495A40; // type:function size:0x14
-__nw__FUlPv = .text:0x80495A60; // type:function size:0x8
-__nw__FUlPQ23EGG4Heapi = .text:0x80495A70; // type:function size:0x10
-__nwa__FUl = .text:0x80495A80; // type:function size:0x10
-__nwa__FUli = .text:0x80495A90; // type:function size:0x8
-__nwa__FUlPQ23EGG4Heapi = .text:0x80495AA0; // type:function size:0x10
+0x80495A60 = .text:0x80495A60; // type:function size:0x8
+0x80495A70 = .text:0x80495A70; // type:function size:0x10
+0x80495A80 = .text:0x80495A80; // type:function size:0x10
+0x80495A90 = .text:0x80495A90; // type:function size:0x8
+0x80495AA0 = .text:0x80495AA0; // type:function size:0x10
 __ct__Q23EGG7ExpHeapFP12MEMiHeapHead = .text:0x80495AB0; // type:function size:0x3C
 __dt__Q23EGG7ExpHeapFv = .text:0x80495AF0; // type:function size:0x74
 create__Q23EGG7ExpHeapFPvUlUs = .text:0x80495B70; // type:function size:0xB8
@@ -30707,7 +30707,7 @@ lbl_8050D37C = .data:0x8050D37C; // type:object size:0x24
 lbl_8050D3A0 = .data:0x8050D3A0; // type:object size:0x10 data:string
 lbl_8050D3B0 = .data:0x8050D3B0; // type:object size:0x14
 lbl_8050D3C4 = .data:0x8050D3C4; // type:object size:0x7C
-lbl_8050D440 = .data:0x8050D440; // type:object size:0x18
+__vt__14dHeapAllocator = .data:0x8050D440; // type:object size:0x18
 lbl_8050D458 = .data:0x8050D458; // type:object size:0xC data:string
 lbl_8050D464 = .data:0x8050D464; // type:object size:0x10
 lbl_8050D474 = .data:0x8050D474; // type:object size:0x10
@@ -39735,7 +39735,7 @@ layoutExHeap__5dHeap = .sbss:0x805751B8; // type:object size:0x4 data:4byte
 layoutEx2Heap__5dHeap = .sbss:0x805751BC; // type:object size:0x4 data:4byte
 layoutResHeap__5dHeap = .sbss:0x805751C0; // type:object size:0x4 data:4byte
 fontHeap__5dHeap = .sbss:0x805751C4; // type:object size:0x4 data:4byte
-HBMHeap__5dHeap = .sbss:0x805751C8; // type:object size:0x8 data:4byte
+HBMHeap__5dHeap = .sbss:0x805751C8; // type:object size:0x4 data:4byte
 lbl_805751D0 = .sbss:0x805751D0; // type:object size:0x8 data:4byte
 lbl_805751D8 = .sbss:0x805751D8; // type:object size:0x4 data:4byte
 lbl_805751DC = .sbss:0x805751DC; // type:object size:0x4 data:4byte
@@ -40883,11 +40883,11 @@ sCurrentHeap__Q23EGG4Heap = .sbss:0x80576740; // type:object size:0x4 data:4byte
 lbl_80576744 = .sbss:0x80576744; // type:object size:0x4 data:4byte
 lbl_80576748 = .sbss:0x80576748; // type:object size:0x4 data:4byte
 lbl_8057674C = .sbss:0x8057674C; // type:object size:0x4 data:4byte
-lbl_80576750 = .sbss:0x80576750; // type:object size:0x4 data:4byte
-lbl_80576754 = .sbss:0x80576754; // type:object size:0x4 data:4byte
+sAllocCallback__Q23EGG4Heap = .sbss:0x80576750; // type:object size:0x4 data:4byte
+sFreeCallback__Q23EGG4Heap = .sbss:0x80576754; // type:object size:0x4 data:4byte
 lbl_80576758 = .sbss:0x80576758; // type:object size:0x4 data:4byte
-lbl_8057675C = .sbss:0x8057675C; // type:object size:0x4 data:4byte
-lbl_80576760 = .sbss:0x80576760; // type:object size:0x4 data:4byte
+sAllocCallbackArg__Q23EGG4Heap = .sbss:0x8057675C; // type:object size:0x4 data:4byte
+sFreeCallbackArg__Q23EGG4Heap = .sbss:0x80576760; // type:object size:0x4 data:4byte
 lbl_80576764 = .sbss:0x80576764; // type:object size:0x4 data:4byte
 lbl_80576768 = .sbss:0x80576768; // type:object size:0x8 data:4byte
 lbl_80576770 = .sbss:0x80576770; // type:object size:0x8 data:4byte
@@ -48129,7 +48129,7 @@ lbl_805975C0 = .bss:0x805975C0; // type:object size:0x48
 lbl_80597608 = .bss:0x80597608; // type:object size:0x10
 lbl_80597618 = .bss:0x80597618; // type:object size:0xA8
 lbl_805976C0 = .bss:0x805976C0; // type:object size:0x80 data:byte
-lbl_80597740 = .bss:0x80597740; // type:object size:0x18 data:4byte
+sAllocator__14dHeapAllocator = .bss:0x80597740; // type:object size:0x18 data:4byte
 lbl_80597758 = .bss:0x80597758; // type:object size:0x318
 lbl_80597A70 = .bss:0x80597A70; // type:object size:0x10
 lbl_80597A80 = .bss:0x80597A80; // type:object size:0x8B70 data:float
diff --git a/include/d/d_heap_alloc.h b/include/d/d_heap_alloc.h
new file mode 100644
index 00000000..726a997d
--- /dev/null
+++ b/include/d/d_heap_alloc.h
@@ -0,0 +1,82 @@
+#ifndef D_HEAP_ALLOC
+#define D_HEAP_ALLOC
+
+#include <egg/core/eggHeap.h>
+
+class dHeapAllocatorBase {
+public:
+    /* vtable at 0x00 */
+    dHeapAllocatorBase()
+        : mCallbacksInitialized(0), mPreviousAllocCallback(nullptr), mPreviousAllocCallbackArg(nullptr),
+          mPreviousFreeCallback(nullptr), mPreviousFreeCallbackArg(nullptr) {}
+    inline void doInitCallbacks() {
+        if (!mCallbacksInitialized) {
+            mCallbacksInitialized = 1;
+
+            void *oldAllocCallbackArg = EGG::Heap::sAllocCallbackArg;
+            void *oldFreeCallbackArg = EGG::Heap::sFreeCallbackArg;
+            EGG::HeapAllocCallback oldAllocCallback = EGG::Heap::sAllocCallback;
+            EGG::HeapFreeCallback oldFreeCallback = EGG::Heap::sFreeCallback;
+
+            EGG::HeapFreeCallback freeThunk = &freeCallback;
+            EGG::HeapAllocCallback allocThunk = &allocCallback;
+
+            EGG::Heap::sAllocCallback = allocThunk;
+            EGG::Heap::sAllocCallbackArg = this;
+            mPreviousAllocCallback = oldAllocCallback;
+            mPreviousAllocCallbackArg = oldAllocCallbackArg;
+
+            EGG::Heap::sFreeCallback = freeThunk;
+            EGG::Heap::sFreeCallbackArg = this;
+            mPreviousFreeCallback = oldFreeCallback;
+            mPreviousFreeCallbackArg = oldFreeCallbackArg;
+        }
+    }
+    static void allocCallback(EGG::HeapAllocArg *arg) {
+        dHeapAllocatorBase *allocator = (dHeapAllocatorBase *)(arg->userArg);
+        allocator->onAlloc(arg);
+        if (allocator->mPreviousAllocCallback) {
+            EGG::HeapAllocArg chainArg = *arg;
+            chainArg.userArg = allocator->mPreviousAllocCallbackArg;
+            (allocator->mPreviousAllocCallback)(&chainArg);
+        }
+    };
+    static void freeCallback(EGG::HeapFreeArg *arg) {
+        dHeapAllocatorBase *allocator = (dHeapAllocatorBase *)(arg->userArg);
+        EGG::HeapFreeArg chainArg;
+        allocator->onFree(arg);
+        if (allocator->mPreviousFreeCallback) {
+            chainArg = *arg;
+            chainArg.userArg = allocator->mPreviousFreeCallbackArg;
+            (allocator->mPreviousFreeCallback)(&chainArg);
+        }
+    }
+
+    virtual ~dHeapAllocatorBase() {}
+    virtual void onAlloc(EGG::HeapAllocArg *arg) {};
+    virtual void onFree(EGG::HeapFreeArg *arg) {};
+
+    /* 0x04 */ bool mCallbacksInitialized;
+    /* 0x08 */ EGG::HeapAllocCallback mPreviousAllocCallback;
+    /* 0x0C */ void *mPreviousAllocCallbackArg;
+    /* 0x10 */ EGG::HeapFreeCallback mPreviousFreeCallback;
+    /* 0x14 */ void *mPreviousFreeCallbackArg;
+};
+
+void *operator new(size_t size);
+void *operator new[](size_t size);
+void operator delete(void *ptr);
+void operator delete[](void *ptr);
+
+class dHeapAllocator : public dHeapAllocatorBase {
+public:
+    dHeapAllocator() {}
+    virtual ~dHeapAllocator() {}
+    virtual void onAlloc(EGG::HeapAllocArg *arg);
+
+    static void initCallbacks();
+
+    static dHeapAllocator sAllocator;
+};
+
+#endif
diff --git a/include/egg/core/eggHeap.h b/include/egg/core/eggHeap.h
index 10b114e9..0bd0b508 100644
--- a/include/egg/core/eggHeap.h
+++ b/include/egg/core/eggHeap.h
@@ -17,14 +17,15 @@ namespace EGG {
 class Allocator;
 
 struct HeapAllocArg {
-    int userArg; // 00
+    void *userArg; // 00
     u32 size;    // 04
     int align;   // 08
     Heap *heap;  // 0C heap to allocate in
+    int another; // 10
 
     inline HeapAllocArg() : userArg(0), size(0), align(0), heap(nullptr) {}
 };
-typedef void (*HeapAllocCallback)(HeapAllocArg &arg);
+typedef void (*HeapAllocCallback)(HeapAllocArg *arg);
 
 struct HeapErrorArg {
     const char *msg;
@@ -32,13 +33,14 @@ struct HeapErrorArg {
 
     inline HeapErrorArg() {}
 };
-typedef void (*ErrorCallback)(void *);
+typedef void (*ErrorCallback)(HeapErrorArg *);
 
 struct HeapFreeArg {
-    u32 arg1; // Idk the args
-    u32 arg2;
+    void *userArg;
+    int arg1;
+    int arg2;
 };
-typedef void (*HeapFreeCallback)(void *);
+typedef void (*HeapFreeCallback)(HeapFreeArg *);
 
 typedef void (*HeapCreateCallback)(void *);
 typedef void (*HeapDestroyCallback)(void *);
@@ -151,14 +153,14 @@ public:
     /* 80576760 */ static void *sFreeCallbackArg;
     /* 80576764 */ static HeapCreateCallback sCreateCallback;
     /* 80576764 */ static HeapDestroyCallback sDestroyCallback;
+
+    /* 80495a60 */ void *operator new(size_t, void *p);
+    /* 80495a70 */ void *operator new(size_t size, EGG::Heap *heap, u32 align);
+    /* 80495a80 */ void *operator new(size_t size, EGG::Allocator *alloc);
+    /* 80495a90 */ void *operator new[](size_t size, u32 align);
+    /* 80495aa0 */ void *operator new[](size_t size, EGG::Heap *heap, int align);
 };
 
 } // namespace EGG
 
-/* 80495a60 */ void *operator new(size_t, void *p);
-/* 80495a70 */ void *operator new(size_t size, EGG::Heap *heap, u32 align);
-/* 80495a80 */ void *operator new(size_t size, EGG::Allocator *alloc);
-/* 80495a90 */ void *operator new[](size_t size, u32 align);
-/* 80495aa0 */ void *operator new[](size_t size, EGG::Heap *heap, int align);
-
 #endif
diff --git a/src/d/d_heap.cpp b/src/d/d_heap.cpp
index 4bfa5e52..2e9c3198 100644
--- a/src/d/d_heap.cpp
+++ b/src/d/d_heap.cpp
@@ -1,4 +1,5 @@
 #include <d/d_heap.h>
+#include <d/d_heap_alloc.h>
 
 dHeap dHeap::work1Heap;
 dHeap dHeap::work2Heap;
@@ -9,6 +10,22 @@ dHeap dHeap::layoutEx2Heap;
 dHeap dHeap::layoutResHeap;
 dHeap dHeap::fontHeap;
 dHeap dHeap::HBMHeap;
+dHeapAllocator dHeapAllocator::sAllocator;
+
+extern u8 lbl_80571C58;
+extern "C" void fn_802de710();
+
+// TODO TU splits?
+
+void dHeapAllocator::onAlloc(EGG::HeapAllocArg *arg) {
+    if (arg->heap != nullptr) {
+        return;
+    }
+    if (lbl_80571C58 == 0) {
+        return;
+    }
+    fn_802de710();
+}
 
 EGG::ExpHeap *dHeap::init(const char *name, size_t size, EGG::Heap *parent) {
     heap = EGG::ExpHeap::create(size, parent, 4);
@@ -18,30 +35,61 @@ EGG::ExpHeap *dHeap::init(const char *name, size_t size, EGG::Heap *parent) {
     }
     return heap;
 }
+
 void dHeap::createWork1Heap(size_t size, EGG::Heap *parent) {
-    work1Heap.init("ゲーム用作業用ヒープ(dHeap::work1Heap)", size, parent);
+    static const char name[] = "ゲーム用作業用ヒープ(dHeap::work1Heap)";
+    work1Heap.init(name, size, parent);
 }
 void dHeap::createWork2Heap(size_t size, EGG::Heap *parent) {
-    work2Heap.init("ゲーム用作業用ヒープ(dHeap::work2Heap)", size, parent);
+    static const char name[] = "ゲーム用作業用ヒープ(dHeap::work2Heap)";
+    work2Heap.init(name, size, parent);
 }
 void dHeap::createWorkExHeap(size_t size, EGG::Heap *parent) {
-    workExHeap.init("ゲーム用拡張作業用ヒープ(dHeap::workExHeap)", size, parent);
+    static const char name[] = "ゲーム用拡張作業用ヒープ(dHeap::workExHeap)";
+    workExHeap.init(name, size, parent);
 }
 void dHeap::createLayoutHeap(size_t size, EGG::Heap *parent) {
-    layoutHeap.init("レイアウト作業用ヒープ(dHeap::layoutHeap)", size, parent);
+    static const char name[] = "レイアウト作業用ヒープ(dHeap::layoutHeap)";
+    layoutHeap.init(name, size, parent);
 }
 void dHeap::createLayoutExHeap(size_t size, EGG::Heap *parent) {
-    layoutExHeap.init("レイアウト拡張作業用ヒープ(dHeap::layoutExHeap)", size, parent);
+    static const char name[] = "レイアウト拡張作業用ヒープ(dHeap::layoutExHeap)";
+    layoutExHeap.init(name, size, parent);
 }
 void dHeap::createLayoutEx2Heap(size_t size, EGG::Heap *parent) {
-    layoutEx2Heap.init("レイアウト拡張作業用ヒープ2(dHeap::layoutEx2Heap)", size, parent);
+    static const char name[] = "レイアウト拡張作業用ヒープ2(dHeap::layoutEx2Heap)";
+    layoutEx2Heap.init(name, size, parent);
 }
 void dHeap::createLayoutResHeap(size_t size, EGG::Heap *parent) {
-    layoutResHeap.init("レイアウトリソース用ヒープ(dHeap::layoutResHeap)", size, parent);
+    static const char name[] = "レイアウトリソース用ヒープ(dHeap::layoutResHeap)";
+    layoutResHeap.init(name, size, parent);
 }
 void dHeap::createFontHeap(size_t size, EGG::Heap *parent) {
-    fontHeap.init("フォント用ヒープ(dHeap::fontHeap)", size, parent);
+    static const char name[] = "フォント用ヒープ(dHeap::fontHeap)";
+    fontHeap.init(name, size, parent);
 }
 void dHeap::createHBMHeap(size_t size, EGG::Heap *parent) {
-    HBMHeap.init("HBM用ヒープ(dHeap::HBMHeap)", size, parent);
+    static const char name[] = "HBM用ヒープ(dHeap::HBMHeap)";
+    HBMHeap.init(name, size, parent);
+}
+
+// TODO this doesn't match (many more stack stores)
+void dHeapAllocator::initCallbacks() {
+    sAllocator.doInitCallbacks();
+}
+
+void *operator new(size_t size) {
+    return EGG::Heap::alloc(size, 0x04, nullptr);
+}
+
+void *operator new[](size_t size) {
+    return EGG::Heap::alloc(size, 0x04, nullptr);
+}
+
+void operator delete(void *ptr) {
+    return EGG::Heap::free(ptr, nullptr);
+}
+
+void operator delete[](void *ptr) {
+    return EGG::Heap::free(ptr, nullptr);
 }