7 files changed, 73 insertions, 15 deletions

diff --git a/src/core/main.cpp b/src/core/main.cpp
index 843f0671..ea88de59 100644
--- a/src/core/main.cpp
+++ b/src/core/main.cpp
@@ -1601,7 +1601,7 @@ void SystemInit()
 	mwInit();
 #endif
 	
-#ifdef GTA_PS2
+#ifdef USE_CUSTOM_ALLOCATOR
 	InitMemoryMgr();
 #endif
 	
diff --git a/src/fakerw/fake.cpp b/src/fakerw/fake.cpp
index 64a37421..39606335 100644
--- a/src/fakerw/fake.cpp
+++ b/src/fakerw/fake.cpp
@@ -16,9 +16,14 @@ using namespace rw;
 RwUInt8 RwObjectGetType(const RwObject *obj) { return obj->type; }
 
 
-void *RwMalloc(size_t size) { return malloc(size); }
-void *RwCalloc(size_t numObj, size_t sizeObj) { return calloc(numObj, sizeObj); }
-void  RwFree(void *mem) { free(mem); }
+void *RwMalloc(size_t size) { return engine->memfuncs.rwmalloc(size, 0); }
+void *RwCalloc(size_t numObj, size_t sizeObj) {
+	void *mem = RwMalloc(numObj*sizeObj);
+	if(mem)
+		memset(mem, 0, numObj*sizeObj);
+	return mem;
+}
+void  RwFree(void *mem) { engine->memfuncs.rwfree(mem); }
 
 
 //RwReal RwV3dNormalize(RwV3d * out, const RwV3d * in);
@@ -536,8 +541,27 @@ RwBool RwRenderStateSet(RwRenderState state, void *value)
 	}
 }
 
+static rw::MemoryFunctions gMemfuncs;
+static void *(*real_malloc)(size_t size);
+static void *(*real_realloc)(void *mem, size_t newSize);
+static void *mallocWrap(size_t sz, uint32 hint) { if(sz == 0) return nil; return real_malloc(sz); }
+static void *reallocWrap(void *p, size_t sz, uint32 hint) { return real_realloc(p, sz); }
+
+
 // WARNING: unused parameters
-RwBool RwEngineInit(RwMemoryFunctions *memFuncs, RwUInt32 initFlags, RwUInt32 resArenaSize) { Engine::init(); return true; }
+RwBool RwEngineInit(RwMemoryFunctions *memFuncs, RwUInt32 initFlags, RwUInt32 resArenaSize) {
+	if(memFuncs){
+		real_malloc = memFuncs->rwmalloc;
+		real_realloc = memFuncs->rwrealloc;
+		gMemfuncs.rwmalloc = mallocWrap;
+		gMemfuncs.rwrealloc = reallocWrap;
+		gMemfuncs.rwfree = memFuncs->rwfree;
+		Engine::init(&gMemfuncs);
+	}else{
+		Engine::init(nil);
+	}
+	return true;
+}
 // TODO: this is platform dependent
 RwBool RwEngineOpen(RwEngineOpenParams *initParams) {
 	static EngineOpenParams openParams;
diff --git a/src/fakerw/rwplcore.h b/src/fakerw/rwplcore.h
index 79c745b6..511f7678 100644
--- a/src/fakerw/rwplcore.h
+++ b/src/fakerw/rwplcore.h
@@ -141,15 +141,15 @@ RwUInt8 RwObjectGetType(const RwObject *obj);
  ***********************************************
  */
 
-struct RwMemoryFunctions;
-/*
+struct RwMemoryFunctions
 {
+	// NB: from RW 3.6 on the allocating functions take
+	// a hint parameter!
 	void *(*rwmalloc)(size_t size);
 	void  (*rwfree)(void *mem);
 	void *(*rwrealloc)(void *mem, size_t newSize);
 	void *(*rwcalloc)(size_t numObj, size_t sizeObj);
 };
-*/
 
 void *RwMalloc(size_t size);
 void  RwFree(void *mem);
diff --git a/src/rw/MemoryHeap.cpp b/src/rw/MemoryHeap.cpp
index d613a708..2a484df4 100644
--- a/src/rw/MemoryHeap.cpp
+++ b/src/rw/MemoryHeap.cpp
@@ -9,8 +9,11 @@
 
 #ifdef USE_CUSTOM_ALLOCATOR
 
-#define MEMORYHEAP_ASSERT(cond) { if (!(cond)) { printf("ASSERT File:%s Line:%d\n", __FILE__, __LINE__); exit(1); } }
-#define MEMORYHEAP_ASSERT_MESSAGE(cond, message) { if (!(cond)) { printf("ASSERT File:%s Line:%d:\n\t%s\n", __FILE__, __LINE__, message); exit(1); } }
+//#define MEMORYHEAP_ASSERT(cond) { if (!(cond)) { printf("ASSERT File:%s Line:%d\n", __FILE__, __LINE__); exit(1); } }
+//#define MEMORYHEAP_ASSERT_MESSAGE(cond, message) { if (!(cond)) { printf("ASSERT File:%s Line:%d:\n\t%s\n", __FILE__, __LINE__, message); exit(1); } }
+
+#define MEMORYHEAP_ASSERT(cond) assert(cond)
+#define MEMORYHEAP_ASSERT_MESSAGE(cond, message) assert(cond)
 
 // registered pointers that we keep track of
 void **gPtrList[4000];
@@ -272,6 +275,7 @@ CMemoryHeap::Free(void *ptr)
 	MEMORYHEAP_ASSERT(m_unkMemId == -1 || m_unkMemId == block->m_memId);
 
 	RegisterFree(block);
+	block->m_memId = MEMID_FREE;
 	CombineFreeBlocks(block);
 	FreeBlock(block);
 	if(block->m_ptrListIndex != -1){
@@ -313,7 +317,7 @@ uint32
 CMemoryHeap::CombineFreeBlocks(HeapBlockDesc *block)
 {
 	HeapBlockDesc *next = block->GetNextConsecutive();
-	if(next->m_memId == MEMID_FREE)
+	if(next->m_memId != MEMID_FREE)
 		return block->m_size;
 	// get rid of free blocks after this one and adjust size
 	for(; next->m_memId == MEMID_FREE; next = next->GetNextConsecutive())
@@ -535,6 +539,10 @@ MemoryMgrCalloc(uint32 num, uint32 size)
 void
 MemoryMgrFree(void *ptr)
 {
+#ifdef FIX_BUGS
+	// i don't suppose this is handled by RW?
+	if(ptr == nil) return;
+#endif
 	gMainHeap.Free(ptr);
 }
 
diff --git a/src/rw/MemoryHeap.h b/src/rw/MemoryHeap.h
index 840e016a..3f6fb5a0 100644
--- a/src/rw/MemoryHeap.h
+++ b/src/rw/MemoryHeap.h
@@ -1,6 +1,12 @@
 #pragma once
 
+// some windows shit
+#ifdef MoveMemory
+#undef MoveMemory
+#endif
+
 extern RwMemoryFunctions memFuncs;
+void InitMemoryMgr(void);
 
 template<typename T, uint32 N>
 class CStack
@@ -56,7 +62,10 @@ struct HeapBlockDesc
 	}
 };
 
+#ifdef USE_CUSTOM_ALLOCATOR
+// TODO: figure something out for 64 bit pointers
 static_assert(sizeof(HeapBlockDesc) == 0x10, "HeapBlockDesc must have 0x10 size otherwise most of assumptions don't make sense");
+#endif
 
 struct HeapBlockList
 {
@@ -181,8 +190,7 @@ public:
 				return;
 			}
 		}
-		HeapBlockDesc *it;
-		for(it = m_freeList.m_first.m_next; it->m_size < block->m_size; it = it->m_next);
-		block->InsertHeapFreeBlock(it->m_prev);
+		HeapBlockDesc *b = m_freeList.m_first.FindSmallestFreeBlock(block->m_size);
+		block->InsertHeapFreeBlock(b->m_prev);
 	}
 };
diff --git a/src/skel/glfw/glfw.cpp b/src/skel/glfw/glfw.cpp
index 982e8641..2722a4df 100644
--- a/src/skel/glfw/glfw.cpp
+++ b/src/skel/glfw/glfw.cpp
@@ -40,6 +40,7 @@
 #include "Sprite2d.h"
 #include "AnimViewer.h"
 #include "Font.h"
+#include "MemoryHeap.h"
 
 #define MAX_SUBSYSTEMS		(16)
 
@@ -277,7 +278,11 @@ psMouseSetPos(RwV2d *pos)
 RwMemoryFunctions*
 psGetMemoryFunctions(void)
 {
+#ifdef USE_CUSTOM_ALLOCATOR
+	return &memFuncs;
+#else
 	return nil;
+#endif
 }
 
 /*
@@ -1461,6 +1466,10 @@ main(int argc, char *argv[])
 	RwV2d pos;
 	RwInt32 i;
 
+#ifdef USE_CUSTOM_ALLOCATOR
+	InitMemoryMgr();
+#endif
+
 #ifndef _WIN32
 	struct sigaction act;
 	act.sa_sigaction = terminateHandler;
diff --git a/src/skel/win/win.cpp b/src/skel/win/win.cpp
index 9effaa31..b4897d67 100644
--- a/src/skel/win/win.cpp
+++ b/src/skel/win/win.cpp
@@ -97,6 +97,7 @@ static psGlobalType PsGlobal;
 #include "Sprite2d.h"
 #include "AnimViewer.h"
 #include "Font.h"
+#include "MemoryHeap.h"
 
 VALIDATE_SIZE(psGlobalType, 0x28);
 
@@ -304,7 +305,11 @@ psMouseSetPos(RwV2d *pos)
 RwMemoryFunctions*
 psGetMemoryFunctions(void)
 {
+#ifdef USE_CUSTOM_ALLOCATOR
+	return &memFuncs;
+#else
 	return nil;
+#endif
 }
 
 /*
@@ -2006,7 +2011,11 @@ WinMain(HINSTANCE instance,
 	RwChar **argv;
 	SystemParametersInfo(SPI_SETFOREGROUNDLOCKTIMEOUT, 0, nil, SPIF_SENDCHANGE);
 
-#if 0
+#ifdef USE_CUSTOM_ALLOCATOR
+	InitMemoryMgr();
+#endif
+
+#if 1
 	// TODO: make this an option somewhere
 	AllocConsole();
 	freopen("CONIN$", "r", stdin);