8 files changed, 803 insertions, 9 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
new file mode 100644
index 00000000..2a484df4
--- /dev/null
+++ b/src/rw/MemoryHeap.cpp
@@ -0,0 +1,556 @@
+#include "common.h"
+#include "main.h"
+#include "FileMgr.h"
+#include "Timer.h"
+#include "ModelInfo.h"
+#include "Streaming.h"
+#include "FileLoader.h"
+#include "MemoryHeap.h"
+
+#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) assert(cond)
+#define MEMORYHEAP_ASSERT_MESSAGE(cond, message) assert(cond)
+
+// registered pointers that we keep track of
+void **gPtrList[4000];
+int32 numPtrs;
+int32 gPosnInList;
+// indices into the ptr list in here are free
+CStack<int32, 4000> m_ptrListIndexStack;
+// how much memory we've moved
+uint32 memMoved;
+
+CMemoryHeap gMainHeap;
+
+void
+CMemoryHeap::Init(uint32 total)
+{
+	MEMORYHEAP_ASSERT((total != 0xF) != 0);
+
+	m_totalMemUsed = 0;
+	m_memUsed = nil;
+	m_currentMemID = MEMID_FREE;
+	m_blocksUsed = nil;
+	m_totalBlocksUsed = 0;
+	m_unkMemId = -1;
+
+	uint8 *mem = (uint8*)malloc(total);
+	assert(((uintptr)mem & 0xF) == 0);
+	m_start = (HeapBlockDesc*)mem;
+	m_end = (HeapBlockDesc*)(mem + total - sizeof(HeapBlockDesc));
+	m_start->m_memId = MEMID_FREE;
+	m_start->m_size = total - 2*sizeof(HeapBlockDesc);
+	m_end->m_memId = MEMID_ID1;
+	m_end->m_size = 0;
+
+	m_freeList.m_last.m_size = INT_MAX;
+	m_freeList.Init();
+	m_freeList.Insert(m_start);
+
+	// TODO: figure out what these are and use sizeof
+	m_fixedSize[0].Init(0x10);
+	m_fixedSize[1].Init(0x20);
+	m_fixedSize[2].Init(0xE0);
+	m_fixedSize[3].Init(0x60);
+	m_fixedSize[4].Init(0x1C0);
+	m_fixedSize[5].Init(0x50);
+
+	m_currentMemID = MEMID_FREE;	// disable registration
+	m_memUsed = (uint32*)Malloc(NUM_MEMIDS * sizeof(uint32));
+	m_blocksUsed = (uint32*)Malloc(NUM_MEMIDS * sizeof(uint32));
+	RegisterMalloc(GetDescFromHeapPointer(m_memUsed));
+	RegisterMalloc(GetDescFromHeapPointer(m_blocksUsed));
+
+	m_currentMemID = MEMID_ID1;
+	for(int i = 0; i < NUM_MEMIDS; i++){
+		m_memUsed[i] = 0;
+		m_blocksUsed[i] = 0;
+	}
+}
+
+void
+CMemoryHeap::RegisterMalloc(HeapBlockDesc *block)
+{
+	block->m_memId = m_currentMemID;
+	if(m_currentMemID == MEMID_FREE)
+		return;
+	m_totalMemUsed += block->m_size + sizeof(HeapBlockDesc);
+	m_memUsed[m_currentMemID] += block->m_size + sizeof(HeapBlockDesc);
+	m_blocksUsed[m_currentMemID]++;
+	m_totalBlocksUsed++;
+}
+
+void
+CMemoryHeap::RegisterFree(HeapBlockDesc *block)
+{
+	if(block->m_memId == MEMID_FREE)
+		return;
+	m_totalMemUsed -= block->m_size + sizeof(HeapBlockDesc);
+	m_memUsed[m_currentMemID] -= block->m_size + sizeof(HeapBlockDesc);
+	m_blocksUsed[m_currentMemID]--;
+	m_totalBlocksUsed--;
+}
+
+void*
+CMemoryHeap::Malloc(uint32 size)
+{
+	static int recursion = 0;
+
+	// weird way to round up
+	if((size & 0xF) != 0)
+		size = (size&~0xF) + 0x10;
+
+	recursion++;
+
+	// See if we can allocate from one of the fixed-size lists
+	for(int i = 0; i < NUM_FIXED_MEMBLOCKS; i++){
+		CommonSize *list = &m_fixedSize[i];
+		if(m_fixedSize[i].m_size == size){
+			HeapBlockDesc *block = list->Malloc();
+			if(block){
+				RegisterMalloc(block);
+				recursion--;
+				return block->GetDataPointer();
+			}
+			break;
+		}
+	}
+
+	// now try the normal free list
+	HeapBlockDesc *next;
+	for(HeapBlockDesc *block = m_freeList.m_first.m_next;
+	    block != &m_freeList.m_last;
+	    block = next){
+		MEMORYHEAP_ASSERT(block->m_memId == MEMID_FREE);
+		MEMORYHEAP_ASSERT_MESSAGE(block >= m_start && block <= m_end, "Block outside of memory");
+
+		// make sure block has maximum size
+		uint32 initialsize = block->m_size;
+		uint32 blocksize = CombineFreeBlocks(block);
+#ifdef FIX_BUGS
+		// has to be done here because block can be moved
+		next = block->m_next;
+#endif
+		if(initialsize != blocksize){
+			block->RemoveHeapFreeBlock();
+			HeapBlockDesc *pos = block->m_prev->FindSmallestFreeBlock(block->m_size);
+			block->InsertHeapFreeBlock(pos->m_prev);
+		}
+		if(block->m_size >= size){
+			// got space to allocate from!
+			block->RemoveHeapFreeBlock();
+			FillInBlockData(block, block->GetNextConsecutive(), size);
+			recursion--;
+			return block->GetDataPointer();
+		}
+#ifndef FIX_BUGS
+		next = block->m_next;
+#endif
+	}
+
+	// oh no, we're losing, try to free some stuff
+	static bool removeCollision = false;
+	static bool removeIslands = false;
+	static bool removeBigBuildings = false;
+	size_t initialMemoryUsed = CStreaming::ms_memoryUsed;
+	CStreaming::MakeSpaceFor(0xCFE800 - CStreaming::ms_memoryUsed);
+	if (recursion > 10)
+		CGame::TidyUpMemory(true, false);
+	else if (recursion > 6)
+		CGame::TidyUpMemory(false, true);
+	if (initialMemoryUsed == CStreaming::ms_memoryUsed && recursion > 11) {
+		if (!removeCollision && !CGame::playingIntro) {
+			CModelInfo::RemoveColModelsFromOtherLevels(LEVEL_GENERIC);
+			removeCollision = true;
+		}
+		else if (!removeIslands && !CGame::playingIntro) {
+			CStreaming::RemoveIslandsNotUsed(LEVEL_INDUSTRIAL);
+			CStreaming::RemoveIslandsNotUsed(LEVEL_COMMERCIAL);
+			CStreaming::RemoveIslandsNotUsed(LEVEL_SUBURBAN);
+			removeIslands = true;
+		}
+		else if (!removeBigBuildings) {
+			CStreaming::RemoveBigBuildings(LEVEL_INDUSTRIAL);
+			CStreaming::RemoveBigBuildings(LEVEL_COMMERCIAL);
+			CStreaming::RemoveBigBuildings(LEVEL_SUBURBAN);
+		}
+		else {
+			LoadingScreen("NO MORE MEMORY", nil, nil);
+			LoadingScreen("NO MORE MEMORY", nil, nil);
+		}
+		CGame::TidyUpMemory(true, false);
+	}
+	void *mem = Malloc(size);
+	if (removeCollision) {
+		CTimer::Stop();
+		// different on PS2
+		CFileLoader::LoadCollisionFromDatFile(CCollision::ms_collisionInMemory);
+		removeCollision = false;
+		CTimer::Update();
+	}
+	if (removeBigBuildings || removeIslands) {
+		CTimer::Stop();
+		if (!CGame::playingIntro)
+			CStreaming::RequestBigBuildings(CGame::currLevel);
+		CStreaming::LoadAllRequestedModels(true);
+		removeBigBuildings = false;
+		removeIslands = false;
+		CTimer::Update();
+	}
+	recursion--;
+	return mem;
+}
+
+void*
+CMemoryHeap::Realloc(void *ptr, uint32 size)
+{
+	if(ptr == nil)
+		return Malloc(size);
+
+	// weird way to round up
+	if((size & 0xF) != 0)
+		size = (size&~0xF) + 0x10;
+
+	HeapBlockDesc *block = GetDescFromHeapPointer(ptr);
+
+#ifdef FIX_BUGS
+	// better handling of size < block->m_size
+	if(size == 0){
+		Free(ptr);
+		return nil;
+	}
+	if(block->m_size >= size){
+		// shrink allocated block
+		RegisterFree(block);
+		PushMemId(block->m_memId);
+		FillInBlockData(block, block->GetNextConsecutive(), size);
+		PopMemId();
+		return ptr;
+	}
+#else
+	// not growing. just returning here is a bit cheap though
+	if(block->m_size >= size)
+		return ptr;
+#endif
+
+	// have to grow allocated block
+	HeapBlockDesc *next = block->GetNextConsecutive();
+	MEMORYHEAP_ASSERT_MESSAGE(next >= m_start && next <= m_end, "Block outside of memory");
+	if(next->m_memId == MEMID_FREE){
+		// try to grow the current block
+		// make sure the next free block has maximum size
+		uint32 freespace = CombineFreeBlocks(next);
+		HeapBlockDesc *end = next->GetNextConsecutive();
+		MEMORYHEAP_ASSERT_MESSAGE(end >= m_start && end <= m_end, "Block outside of memory");
+		// why the sizeof here?
+		if(block->m_size + next->m_size + sizeof(HeapBlockDesc) >= size){
+			// enough space to grow
+			next->RemoveHeapFreeBlock();
+			RegisterFree(block);
+			PushMemId(block->m_memId);
+			FillInBlockData(block, next->GetNextConsecutive(), size);
+			PopMemId();
+			return ptr;
+		}
+	}
+
+	// can't grow the existing block, have to get a new one and copy
+	PushMemId(block->m_memId);
+	void *dst = Malloc(size);
+	PopMemId();
+	memcpy(dst, ptr, block->m_size);
+	Free(ptr);
+	return dst;
+}
+
+void
+CMemoryHeap::Free(void *ptr)
+{
+	HeapBlockDesc *block = GetDescFromHeapPointer(ptr);
+	MEMORYHEAP_ASSERT_MESSAGE(block->m_memId != MEMID_FREE, "MemoryHeap corrupt");
+	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){
+		int32 idx = block->m_ptrListIndex;
+		gPtrList[idx] = nil;
+		m_ptrListIndexStack.push(idx);
+	}
+	block->m_ptrListIndex = -1;
+}
+
+// allocate 'size' bytes from 'block'
+void
+CMemoryHeap::FillInBlockData(HeapBlockDesc *block, HeapBlockDesc *end, uint32 size)
+{
+	block->m_size = size;
+	block->m_ptrListIndex = -1;
+	HeapBlockDesc *remainder = block->GetNextConsecutive();
+	MEMORYHEAP_ASSERT(remainder <= end);
+
+	if(remainder < end-1){
+		RegisterMalloc(block);
+
+		// can fit another block in the remaining space
+		remainder->m_size = GetSizeBetweenBlocks(remainder, end);
+		remainder->m_memId = MEMID_FREE;
+		MEMORYHEAP_ASSERT(remainder->m_size != 0);
+		FreeBlock(remainder);
+	}else{
+		// fully allocate this one
+		if(remainder < end)
+			// no gaps allowed
+			block->m_size = GetSizeBetweenBlocks(block, end);
+		RegisterMalloc(block);
+	}
+}
+
+// Make sure free block has no other free blocks after it
+uint32
+CMemoryHeap::CombineFreeBlocks(HeapBlockDesc *block)
+{
+	HeapBlockDesc *next = block->GetNextConsecutive();
+	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())
+		next->RemoveHeapFreeBlock();
+	block->m_size = GetSizeBetweenBlocks(block, next);
+	return block->m_size;
+}
+
+// Try to move all registered memory blocks into more optimal location
+void
+CMemoryHeap::TidyHeap(void)
+{
+	for(int i = 0; i < numPtrs; i++){
+		if(gPtrList[i] == nil || *gPtrList[i] == nil)
+			continue;
+		HeapBlockDesc *newblock = WhereShouldMemoryMove(*gPtrList[i]);
+		if(newblock)
+			*gPtrList[i] = MoveHeapBlock(newblock, GetDescFromHeapPointer(*gPtrList[i]));
+	}
+}
+
+// 
+void
+CMemoryHeap::RegisterMemPointer(void *ptr)
+{
+	HeapBlockDesc *block = GetDescFromHeapPointer(*(void**)ptr);
+
+	if(block->m_ptrListIndex != -1)
+		return;		// already registered
+
+	int index;
+	if(m_ptrListIndexStack.sp > 0){
+		// re-use a previously free'd index
+		index = m_ptrListIndexStack.pop();
+	}else{
+		// have to find a new index
+		index = gPosnInList;
+
+		void **pp = gPtrList[index];
+		// we're replacing an old pointer here??
+		if(pp && *pp && *pp != (void*)0xDDDDDDDD)
+			GetDescFromHeapPointer(*pp)->m_ptrListIndex = -1;
+
+		gPosnInList++;
+		if(gPosnInList == 4000)
+			gPosnInList = 0;
+		if(numPtrs < 4000)
+			numPtrs++;
+	}
+	gPtrList[index] = (void**)ptr;
+	block->m_ptrListIndex = index;
+}
+
+void*
+CMemoryHeap::MoveMemory(void *ptr)
+{
+	HeapBlockDesc *newblock = WhereShouldMemoryMove(ptr);
+	if(newblock)
+		return MoveHeapBlock(newblock, GetDescFromHeapPointer(ptr));
+	else
+		return ptr;
+}
+
+HeapBlockDesc*
+CMemoryHeap::WhereShouldMemoryMove(void *ptr)
+{
+	HeapBlockDesc *block = GetDescFromHeapPointer(ptr);
+	MEMORYHEAP_ASSERT(block->m_memId != MEMID_FREE);
+
+	HeapBlockDesc *next = block->GetNextConsecutive();
+	if(next->m_memId != MEMID_FREE)
+		return nil;
+
+	// we want to move the block into another block
+	// such that the free space between this and the next block can be minimized
+	HeapBlockDesc *newblock = m_freeList.m_first.FindSmallestFreeBlock(block->m_size);
+	// size of free space wouldn't decrease, so return
+	if(newblock->m_size >= block->m_size + next->m_size)
+		return nil;
+	// size of free space wouldn't decrease enough
+	if(newblock->m_size >= 16 + 1.125f*block->m_size)	// what are 16 and 1.125 here? sizeof(HeapBlockDesc)?
+		return nil;
+	return newblock;
+}
+
+void*
+CMemoryHeap::MoveHeapBlock(HeapBlockDesc *dst, HeapBlockDesc *src)
+{
+	PushMemId(src->m_memId);
+	dst->RemoveHeapFreeBlock();
+	FillInBlockData(dst, dst->GetNextConsecutive(), src->m_size);
+	PopMemId();
+	memcpy(dst->GetDataPointer(), src->GetDataPointer(), src->m_size);
+	memMoved += src->m_size;
+	dst->m_ptrListIndex = src->m_ptrListIndex;
+	src->m_ptrListIndex = -1;
+	Free(src->GetDataPointer());
+	return dst->GetDataPointer();
+}
+
+uint32
+CMemoryHeap::GetMemoryUsed(int32 id)
+{
+	return m_memUsed[id];
+}
+
+uint32
+CMemoryHeap::GetBlocksUsed(int32 id)
+{
+	return m_blocksUsed[id];
+}
+
+void
+CMemoryHeap::PopMemId(void)
+{
+	m_currentMemID = m_idStack.pop();
+}
+
+void
+CMemoryHeap::PushMemId(int32 id)
+{
+	MEMORYHEAP_ASSERT(id != MEMID_FREE);
+	m_idStack.push(m_currentMemID);
+	m_currentMemID = id;
+}
+
+void
+CMemoryHeap::ParseHeap(void)
+{
+	char tmp[16];
+	int fd = CFileMgr::OpenFileForWriting("heap.txt");
+	CTimer::Stop();
+
+	// CMemoryHeap::IntegrityCheck();
+
+	uint32 addrQW = 0;
+	for(HeapBlockDesc *block = m_start; block < m_end; block = block->GetNextConsecutive()){
+		char chr = '*';	// free
+		if(block->m_memId != MEMID_FREE)
+			chr = block->m_memId-MEMID_ID1 + 'A';
+		int numQW = block->m_size>>4;
+
+		if((addrQW & 0x3F) == 0){
+			sprintf(tmp, "\n%5dK:", addrQW>>6);
+			CFileMgr::Write(fd, tmp, 8);
+		}
+		CFileMgr::Write(fd, "#", 1);	// the descriptor, has to be 16 bytes!!!!
+		addrQW++;
+
+		while(numQW--){
+			if((addrQW & 0x3F) == 0){
+				sprintf(tmp, "\n%5dK:", addrQW>>6);
+				CFileMgr::Write(fd, tmp, 8);
+			}
+			CFileMgr::Write(fd, &chr, 1);
+			addrQW++;
+		}
+	}
+
+	CTimer::Update();
+	CFileMgr::CloseFile(fd);
+}
+
+
+void
+CommonSize::Init(uint32 size)
+{
+	m_freeList.Init();
+	m_size = size;
+	m_failed = 0;
+	m_remaining = 0;
+}
+
+
+
+void *pMemoryTop;
+
+void
+InitMemoryMgr(void)
+{
+#ifdef GTA_PS2
+#error "finish this"
+#else
+	// randomly allocate 128mb
+	gMainHeap.Init(128*1024*1024);
+#endif
+}
+
+void*
+MemoryMgrMalloc(uint32 size)
+{
+	void *mem = gMainHeap.Malloc(size);
+	if(mem > pMemoryTop)
+		pMemoryTop = mem;
+	return mem;
+}
+
+void*
+MemoryMgrRealloc(void *ptr, uint32 size)
+{
+	void *mem = gMainHeap.Realloc(ptr, size);
+	if(mem > pMemoryTop)
+		pMemoryTop = mem;
+	return mem;
+}
+
+void*
+MemoryMgrCalloc(uint32 num, uint32 size)
+{
+	void *mem = gMainHeap.Malloc(num*size);
+	if(mem > pMemoryTop)
+		pMemoryTop = mem;
+#ifdef FIX_BUGS
+	memset(mem, 0, num*size);
+#endif
+	return mem;
+}
+
+void
+MemoryMgrFree(void *ptr)
+{
+#ifdef FIX_BUGS
+	// i don't suppose this is handled by RW?
+	if(ptr == nil) return;
+#endif
+	gMainHeap.Free(ptr);
+}
+
+RwMemoryFunctions memFuncs = {
+	MemoryMgrMalloc,
+	MemoryMgrFree,
+	MemoryMgrRealloc,
+	MemoryMgrCalloc
+};
+
+#endif
diff --git a/src/rw/MemoryHeap.h b/src/rw/MemoryHeap.h
new file mode 100644
index 00000000..3f6fb5a0
--- /dev/null
+++ b/src/rw/MemoryHeap.h
@@ -0,0 +1,196 @@
+#pragma once
+
+// some windows shit
+#ifdef MoveMemory
+#undef MoveMemory
+#endif
+
+extern RwMemoryFunctions memFuncs;
+void InitMemoryMgr(void);
+
+template<typename T, uint32 N>
+class CStack
+{
+public:
+	T values[N];
+	uint32 sp;
+
+	CStack() : sp(0) {}
+	void push(const T& val) { values[sp++] = val; }
+	T& pop() { return values[sp--]; }
+};
+
+
+struct HeapBlockDesc
+{
+	uint32 m_size;
+	int16 m_memId;
+	int16 m_ptrListIndex;
+	HeapBlockDesc *m_next;
+	HeapBlockDesc *m_prev;
+
+	HeapBlockDesc *GetNextConsecutive(void)
+	{
+		return (HeapBlockDesc*)((uintptr)this + sizeof(HeapBlockDesc) + m_size);
+	}
+
+	void *GetDataPointer(void)
+	{
+		return (void*)((uintptr)this + sizeof(HeapBlockDesc));
+	}
+
+	void RemoveHeapFreeBlock(void)
+	{
+		m_next->m_prev = m_prev;
+		m_prev->m_next = m_next;
+	}
+
+	// after node
+	void InsertHeapFreeBlock(HeapBlockDesc *node)
+	{
+		m_next = node->m_next;
+		node->m_next->m_prev = this;
+		m_prev = node;
+		node->m_next = this;
+	}
+
+	HeapBlockDesc *FindSmallestFreeBlock(uint32 size)
+	{
+		HeapBlockDesc *b;
+		for(b = m_next; b->m_size < size; b = b->m_next);
+		return b;
+	}
+};
+
+#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
+{
+	HeapBlockDesc m_first;
+	HeapBlockDesc m_last;
+
+	void Init(void)
+	{
+		m_first.m_next = &m_last;
+		m_last.m_prev = &m_first;
+	}
+
+	void Insert(HeapBlockDesc *node)
+	{
+		node->InsertHeapFreeBlock(&m_first);
+	}
+};
+
+struct CommonSize
+{
+	HeapBlockList m_freeList;
+	uint32 m_size;
+	uint32 m_failed;
+	uint32 m_remaining;
+
+	void Init(uint32 size);
+	void Free(HeapBlockDesc *node)
+	{
+		m_freeList.Insert(node);
+		m_remaining++;
+	}
+	HeapBlockDesc *Malloc(void)
+	{
+		if(m_freeList.m_first.m_next == &m_freeList.m_last){
+			m_failed++;
+			return nil;
+		}
+		HeapBlockDesc *block = m_freeList.m_first.m_next;
+		m_remaining--;
+		block->RemoveHeapFreeBlock();
+		block->m_ptrListIndex = -1;
+		return block;
+	}
+};
+
+enum {
+	MEMID_FREE,
+	// IDs from LCS:
+	MEMID_ID1,	// "Game"
+	MEMID_ID2,	// "World"
+	MEMID_ID3,	// "Animation"
+	MEMID_ID4,	// "Pools"
+	MEMID_ID5,	// "Default Models"
+	MEMID_ID6,	// "Streaming"
+	MEMID_ID7,	// "Streamed Models"
+	MEMID_ID8,	// "Streamed LODs"
+	MEMID_ID9,	// "Streamed Textures"
+	MEMID_ID10,	// "Streamed Collision"
+	MEMID_ID11,	// "Streamed Animation"
+	MEMID_ID12,	// "Textures"
+	MEMID_ID13,	// "Collision"
+	MEMID_ID14,	// "PreAlloc"
+	MEMID_ID15,	// "Game Process"
+	MEMID_ID16,	// "Script"
+	MEMID_ID17,	// "Cars"
+	MEMID_ID18,	// "Render"
+	MEMID_ID19,	// "Ped Attr"
+
+	NUM_MEMIDS,
+
+	NUM_FIXED_MEMBLOCKS = 6
+};
+
+class CMemoryHeap
+{
+public:
+	HeapBlockDesc *m_start;
+	HeapBlockDesc *m_end;
+	HeapBlockList m_freeList;
+	CommonSize m_fixedSize[NUM_FIXED_MEMBLOCKS];
+	uint32 m_totalMemUsed;
+	CStack<int32, 16> m_idStack;
+	uint32 m_currentMemID;
+	uint32 *m_memUsed;
+	uint32 m_totalBlocksUsed;
+	uint32 *m_blocksUsed;
+	uint32 m_unkMemId;
+
+	CMemoryHeap(void) : m_start(nil) {}
+	void Init(uint32 total);
+	void RegisterMalloc(HeapBlockDesc *block);
+	void RegisterFree(HeapBlockDesc *block);
+	void *Malloc(uint32 size);
+	void *Realloc(void *ptr, uint32 size);
+	void Free(void *ptr);
+	void FillInBlockData(HeapBlockDesc *block, HeapBlockDesc *end, uint32 size);
+	uint32 CombineFreeBlocks(HeapBlockDesc *block);
+	void *MoveMemory(void *ptr);
+	HeapBlockDesc *WhereShouldMemoryMove(void *ptr);
+	void *MoveHeapBlock(HeapBlockDesc *dst, HeapBlockDesc *src);
+	void PopMemId(void);
+	void PushMemId(int32 id);
+	void RegisterMemPointer(void *ptr);
+	void TidyHeap(void);
+	uint32 GetMemoryUsed(int32 id);
+	uint32 GetBlocksUsed(int32 id);
+
+	void ParseHeap(void);
+
+	HeapBlockDesc *GetDescFromHeapPointer(void *block)
+	{
+		return (HeapBlockDesc*)((uintptr)block - sizeof(HeapBlockDesc));
+	}
+	uint32 GetSizeBetweenBlocks(HeapBlockDesc *first, HeapBlockDesc *second)
+	{
+		return (uintptr)second - (uintptr)first - sizeof(HeapBlockDesc);
+	}
+	void FreeBlock(HeapBlockDesc *block){
+		for(int i = 0; i < NUM_FIXED_MEMBLOCKS; i++){
+			if(m_fixedSize[i].m_size == block->m_size){
+				m_fixedSize[i].Free(block);
+				return;
+			}
+		}
+		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);
diff --git a/vendor/librw b/vendor/librw
-Subproject d9def88c46a742c6bc74bf79021c0f8838480df
+Subproject e8990d5b3d50be72594f93dcc42d749f2976151