dHeap(Allocator) work

2 files changed, 96 insertions, 12 deletions

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