d_heap OK (#47)

* d_heap OK

* wibo workaround

2 files changed, 40 insertions, 16 deletions

diff --git a/include/d/d_heap_alloc.h b/include/d/d_heap_alloc.h
index cbf43c5a..ae7e98e2 100644
--- a/include/d/d_heap_alloc.h
+++ b/include/d/d_heap_alloc.h
@@ -12,21 +12,13 @@ public:
     inline void doInitCallbacks() {
         if (!mCallbacksInitialized) {
             mCallbacksInitialized = 1;
+            EGG::Heap::AllocCallbackBundle prev = EGG::Heap::setAllocCallback(&allocCallback, this);
+            mPreviousAllocCallback = prev.AllocCallback;
+            mPreviousAllocCallbackArg = prev.AllocCallbackArg;
 
-            void *oldAllocCallbackArg = EGG::Heap::sAllocCallbackArg;
-            void *oldFreeCallbackArg = EGG::Heap::sFreeCallbackArg;
-            EGG::HeapAllocCallback oldAllocCallback = EGG::Heap::sAllocCallback;
-            EGG::HeapFreeCallback oldFreeCallback = EGG::Heap::sFreeCallback;
-
-            EGG::Heap::sAllocCallback = &allocCallback;
-            EGG::Heap::sAllocCallbackArg = this;
-            mPreviousAllocCallback = oldAllocCallback;
-            mPreviousAllocCallbackArg = oldAllocCallbackArg;
-
-            EGG::Heap::sFreeCallback = &freeCallback;
-            EGG::Heap::sFreeCallbackArg = this;
-            mPreviousFreeCallback = oldFreeCallback;
-            mPreviousFreeCallbackArg = oldFreeCallbackArg;
+            EGG::Heap::FreeCallbackBundle prev2 = EGG::Heap::setFreeCallback(&freeCallback, this);
+            mPreviousFreeCallback = prev2.FreeCallback;
+            mPreviousFreeCallbackArg = prev2.FreeCallbackArg;
         }
     }
     static void allocCallback(EGG::HeapAllocArg *arg) {
@@ -50,8 +42,8 @@ public:
     }
 
     virtual ~dHeapAllocatorBase() {}
-    virtual void onAlloc(EGG::HeapAllocArg *arg){};
-    virtual void onFree(EGG::HeapFreeArg *arg){};
+    virtual void onAlloc(EGG::HeapAllocArg *arg) {};
+    virtual void onFree(EGG::HeapFreeArg *arg) {};
 
     /* 0x04 */ bool mCallbacksInitialized;
     /* 0x08 */ EGG::HeapAllocCallback mPreviousAllocCallback;
diff --git a/include/egg/core/eggHeap.h b/include/egg/core/eggHeap.h
index e0730703..956a8bfd 100644
--- a/include/egg/core/eggHeap.h
+++ b/include/egg/core/eggHeap.h
@@ -145,6 +145,38 @@ public:
         return mName;
     }
 
+    struct AllocCallbackBundle {
+        HeapAllocCallback AllocCallback;
+        void *AllocCallbackArg;
+    };
+
+    struct FreeCallbackBundle {
+        HeapFreeCallback FreeCallback;
+        void *FreeCallbackArg;
+    };
+
+    static inline AllocCallbackBundle setAllocCallback(HeapAllocCallback AllocCallback, void *AllocCallbackArg) {
+        AllocCallbackBundle ret;
+
+        ret.AllocCallback = sAllocCallback;
+        ret.AllocCallbackArg = sAllocCallbackArg;
+        sAllocCallback = AllocCallback;
+        sAllocCallbackArg = AllocCallbackArg;
+
+        return ret;
+    }
+
+    static inline FreeCallbackBundle setFreeCallback(HeapFreeCallback FreeCallback, void *FreeCallbackArg) {
+        FreeCallbackBundle ret;
+
+        ret.FreeCallback = sFreeCallback;
+        ret.FreeCallbackArg = sFreeCallbackArg;
+        sFreeCallback = FreeCallback;
+        sFreeCallbackArg = FreeCallbackArg;
+
+        return ret;
+    }
+
     /* 80673ae8 */ static nw4r::ut::List sHeapList;
     /* 80673af8 */ static OSMutex sRootMutex;
     /* 80576740 */ static Heap *sCurrentHeap;