5 files changed, 67 insertions, 5 deletions

diff --git a/src/core/FileLoader.cpp b/src/core/FileLoader.cpp
index 88a99fa9..aeaede56 100644
--- a/src/core/FileLoader.cpp
+++ b/src/core/FileLoader.cpp
@@ -243,6 +243,7 @@ CFileLoader::LoadCollisionModel(uint8 *buf, CColModel &model, char *modelname)
 	buf += 44;
 	if(model.numSpheres > 0){
 		model.spheres = (CColSphere*)RwMalloc(model.numSpheres*sizeof(CColSphere));
+		REGISTER_MEMPTR(&model.spheres);
 		for(i = 0; i < model.numSpheres; i++){
 			model.spheres[i].Set(*(float*)buf, *(CVector*)(buf+4), buf[16], buf[17]);
 			buf += 20;
@@ -254,6 +255,7 @@ CFileLoader::LoadCollisionModel(uint8 *buf, CColModel &model, char *modelname)
 	buf += 4;
 	if(model.numLines > 0){
 		model.lines = (CColLine*)RwMalloc(model.numLines*sizeof(CColLine));
+		REGISTER_MEMPTR(&model.lines);
 		for(i = 0; i < model.numLines; i++){
 			model.lines[i].Set(*(CVector*)buf, *(CVector*)(buf+12));
 			buf += 24;
@@ -265,6 +267,7 @@ CFileLoader::LoadCollisionModel(uint8 *buf, CColModel &model, char *modelname)
 	buf += 4;
 	if(model.numBoxes > 0){
 		model.boxes = (CColBox*)RwMalloc(model.numBoxes*sizeof(CColBox));
+		REGISTER_MEMPTR(&model.boxes);
 		for(i = 0; i < model.numBoxes; i++){
 			model.boxes[i].Set(*(CVector*)buf, *(CVector*)(buf+12), buf[24], buf[25]);
 			buf += 28;
@@ -276,6 +279,7 @@ CFileLoader::LoadCollisionModel(uint8 *buf, CColModel &model, char *modelname)
 	buf += 4;
 	if(numVertices > 0){
 		model.vertices = (CompressedVector*)RwMalloc(numVertices*sizeof(CompressedVector));
+		REGISTER_MEMPTR(&model.vertices);
 		for(i = 0; i < numVertices; i++){
 			model.vertices[i].Set(*(float*)buf, *(float*)(buf+4), *(float*)(buf+8));
 			if(Abs(*(float*)buf) >= 256.0f ||
@@ -291,6 +295,7 @@ CFileLoader::LoadCollisionModel(uint8 *buf, CColModel &model, char *modelname)
 	buf += 4;
 	if(model.numTriangles > 0){
 		model.triangles = (CColTriangle*)RwMalloc(model.numTriangles*sizeof(CColTriangle));
+		REGISTER_MEMPTR(&model.triangles);
 		for(i = 0; i < model.numTriangles; i++){
 			model.triangles[i].Set(model.vertices, *(int32*)buf, *(int32*)(buf+4), *(int32*)(buf+8), buf[12], buf[13]);
 			buf += 16;
diff --git a/src/core/Streaming.cpp b/src/core/Streaming.cpp
index a28fe39d..03b49fd6 100644
--- a/src/core/Streaming.cpp
+++ b/src/core/Streaming.cpp
@@ -458,6 +458,35 @@ CStreaming::LoadCdDirectory(const char *dirname, int n)
 	CFileMgr::CloseFile(fd);
 }
 
+#ifdef USE_CUSTOM_ALLOCATOR
+RpAtomic*
+RegisterAtomicMemPtrsCB(RpAtomic *atomic, void *data)
+{
+#if THIS_IS_COMPATIBLE_WITH_GTA3_RW31
+	// not quite sure what's going on here:
+	// gta3's RW 3.1 allocates separate memory for geometry data of RpGeometry.
+	// Is that a R* change? rpDefaultGeometryInstance also depends on it
+	RpGeometry *geo = RpAtomicGetGeometry(atomic);
+	if(geo->triangles)
+		REGISTER_MEMPTR(&geo->triangles);
+	if(geo->matList.materials)
+		REGISTER_MEMPTR(&geo->matList.materials);
+	if(geo->preLitLum)
+		REGISTER_MEMPTR(&geo->preLitLum);
+	if(geo->texCoords[0])
+		REGISTER_MEMPTR(&geo->texCoords[0]);
+	if(geo->texCoords[1])
+		REGISTER_MEMPTR(&geo->texCoords[1]);
+#else
+	// normally RpGeometry is allocated in one block (excluding morph targets)
+	// so we don't really have allocated pointers in the struct.
+	// NB: in librw we actually do it in two allocations (geometry itself and data)
+	// so we could conceivably come up with something here
+#endif
+	return atomic;
+}
+#endif
+
 bool
 CStreaming::ConvertBufferToObject(int8 *buf, int32 streamId)
 {
@@ -494,9 +523,11 @@ CStreaming::ConvertBufferToObject(int8 *buf, int32 streamId)
 
 		PUSH_MEMID(MEMID_STREAM_MODELS);
 		CTxdStore::SetCurrentTxd(mi->GetTxdSlot());
-// TODO(USE_CUSTOM_ALLOCATOR): register mem pointers
 		if(mi->IsSimple()){
 			success = CFileLoader::LoadAtomicFile(stream, streamId);
+#ifdef USE_CUSTOM_ALLOCATOR
+			RegisterAtomicMemPtrsCB(((CSimpleModelInfo*)mi)->m_atomics[0], nil);
+#endif
 		} else if (mi->GetModelType() == MITYPE_VEHICLE) {
 			// load vehicles in two parts
 			CModelInfo::GetModelInfo(streamId)->AddRef();
@@ -505,6 +536,10 @@ CStreaming::ConvertBufferToObject(int8 *buf, int32 streamId)
 				ms_aInfoForModel[streamId].m_loadState = STREAMSTATE_STARTED;
 		}else{
 			success = CFileLoader::LoadClumpFile(stream, streamId);
+#ifdef USE_CUSTOM_ALLOCATOR
+			if(success)
+				RpClumpForAllAtomics((RpClump*)mi->GetRwObject(), RegisterAtomicMemPtrsCB, nil);
+#endif
 		}
 		POP_MEMID();
 		UpdateMemoryUsed();
@@ -628,13 +663,16 @@ CStreaming::FinishLoadingLargeFile(int8 *buf, int32 streamId)
 
 	if(streamId < STREAM_OFFSET_TXD){
 		// Model
-// TODO(USE_CUSTOM_ALLOCATOR): register pointers
 		mi = CModelInfo::GetModelInfo(streamId);
 		PUSH_MEMID(MEMID_STREAM_MODELS);
 		CTxdStore::SetCurrentTxd(mi->GetTxdSlot());
 		success = CFileLoader::FinishLoadClumpFile(stream, streamId);
-		if(success)
+		if(success){
+#ifdef USE_CUSTOM_ALLOCATOR
+			RpClumpForAllAtomics((RpClump*)mi->GetRwObject(), RegisterAtomicMemPtrsCB, nil);
+#endif
 			success = AddToLoadedVehiclesList(streamId);
+		}
 		POP_MEMID();
 		mi->RemoveRef();
 		CTxdStore::RemoveRefWithoutDelete(mi->GetTxdSlot());
diff --git a/src/core/main.cpp b/src/core/main.cpp
index d34eb8f3..102548b6 100644
--- a/src/core/main.cpp
+++ b/src/core/main.cpp
@@ -110,6 +110,9 @@ void TheGame(void);
 void DebugMenuPopulate(void);
 #endif
 
+#ifndef FINAL
+bool gbPrintMemoryUsage;
+#endif
 
 #ifdef GTA_PS2
 #define WANT_TO_LOAD TheMemoryCard.m_bWantToLoad
@@ -957,9 +960,10 @@ DisplayGameDebugText()
 		TWEAKBOOL(bDisplayPosn);
 		TWEAKBOOL(bDisplayRate);
 	}
-#endif
 
-//	PrintMemoryUsage();	// TODO: put this somewhere else
+	if(gbPrintMemoryUsage)
+		PrintMemoryUsage();
+#endif
 
 	char str[200];
 	wchar ustr[200];
diff --git a/src/core/main.h b/src/core/main.h
index 13fff447..77fac46a 100644
--- a/src/core/main.h
+++ b/src/core/main.h
@@ -20,6 +20,10 @@ extern bool gbShowTimebars;
 #define gbShowTimebars false
 #endif
 
+#ifndef FINAL
+extern bool gbPrintMemoryUsage;
+#endif
+
 class CSprite2d;
 
 bool DoRWStuffStartOfFrame(int16 TopRed, int16 TopGreen, int16 TopBlue, int16 BottomRed, int16 BottomGreen, int16 BottomBlue, int16 Alpha);
diff --git a/src/core/re3.cpp b/src/core/re3.cpp
index ee747218..506b2714 100644
--- a/src/core/re3.cpp
+++ b/src/core/re3.cpp
@@ -29,6 +29,7 @@
 #include "Script.h"
 #include "postfx.h"
 #include "custompipes.h"
+#include "MemoryHeap.h"
 
 #ifdef DONT_TRUST_RECOGNIZED_JOYSTICKS
 #include "FileMgr.h"
@@ -383,6 +384,10 @@ SwitchToMission(void)
 }
 #endif
 
+#ifdef USE_CUSTOM_ALLOCATOR
+static void ParseHeap(void) { gMainHeap.ParseHeap(); }
+#endif
+
 static const char *carnames[] = {
 	"landstal", "idaho", "stinger", "linerun", "peren", "sentinel", "patriot", "firetruk", "trash", "stretch", "manana", "infernus", "blista", "pony",
 	"mule", "cheetah", "ambulan", "fbicar", "moonbeam", "esperant", "taxi", "kuruma", "bobcat", "mrwhoop", "bfinject", "corpse", "police", "enforcer",
@@ -565,6 +570,12 @@ DebugMenuPopulate(void)
 		DebugMenuAddVarBool8("Render", "Don't render Objects", &gbDontRenderObjects, nil);
 		DebugMenuAddVarBool8("Render", "Don't Render Water", &gbDontRenderWater, nil);
 
+#ifndef FINAL
+		DebugMenuAddVarBool8("Debug", "Print Memory Usage", &gbPrintMemoryUsage, nil);
+#ifdef USE_CUSTOM_ALLOCATOR
+		DebugMenuAddCmd("Debug", "Parse Heap", ParseHeap);
+#endif
+#endif
 		DebugMenuAddVarBool8("Debug", "Show cullzone debug stuff", &gbShowCullZoneDebugStuff, nil);
 		DebugMenuAddVarBool8("Debug", "Disable zone cull", &gbDisableZoneCull, nil);