1 files changed, 35 insertions, 6 deletions

diff --git a/src/core/Streaming.cpp b/src/core/Streaming.cpp
index e9a7af88..8f1a8891 100644
--- a/src/core/Streaming.cpp
+++ b/src/core/Streaming.cpp
@@ -33,6 +33,7 @@
 #endif
 #include "main.h"
 #include "Frontend.h"
+#include "MemoryHeap.h"
 
 bool CStreaming::ms_disableStreaming;
 bool CStreaming::ms_bLoadingBigModel;
@@ -470,8 +471,10 @@ CStreaming::ConvertBufferToObject(int8 *buf, int32 streamId)
 
 		// Set Txd to use
 		CTxdStore::AddRef(mi->GetTxdSlot());
-		CTxdStore::SetCurrentTxd(mi->GetTxdSlot());
 
+		PUSH_MEMID(MEMID_STREAM_MODELS);
+		CTxdStore::SetCurrentTxd(mi->GetTxdSlot());
+// TODO(USE_CUSTOM_ALLOCATOR): register mem pointers
 		if(mi->IsSimple()){
 			success = CFileLoader::LoadAtomicFile(stream, streamId);
 		} else if (mi->GetModelType() == MITYPE_VEHICLE) {
@@ -483,6 +486,7 @@ CStreaming::ConvertBufferToObject(int8 *buf, int32 streamId)
 		}else{
 			success = CFileLoader::LoadClumpFile(stream, streamId);
 		}
+		POP_MEMID();
 		UpdateMemoryUsed();
 
 		// Txd no longer needed unless we only read part of the file
@@ -506,12 +510,14 @@ CStreaming::ConvertBufferToObject(int8 *buf, int32 streamId)
 			return false;
 		}
 
+		PUSH_MEMID(MEMID_STREAM_TEXUTRES);
 		if(ms_bLoadingBigModel || cdsize > 200){
 			success = CTxdStore::StartLoadTxd(streamId - STREAM_OFFSET_TXD, stream);
 			if(success)
 				ms_aInfoForModel[streamId].m_loadState = STREAMSTATE_STARTED;
 		}else
 		        success = CTxdStore::LoadTxd(streamId - STREAM_OFFSET_TXD, stream);
+		POP_MEMID();
 		UpdateMemoryUsed();
 
 		if(!success){
@@ -561,7 +567,9 @@ CStreaming::ConvertBufferToObject(int8 *buf, int32 streamId)
 	// Mark objects as loaded
 	if(ms_aInfoForModel[streamId].m_loadState != STREAMSTATE_STARTED){
 		ms_aInfoForModel[streamId].m_loadState = STREAMSTATE_LOADED;
+#ifndef USE_CUSTOM_ALLOCATOR
 		ms_memoryUsed += ms_aInfoForModel[streamId].GetCdSize() * CDSTREAM_SECTOR_SIZE;
+#endif
 	}
 
 	endTime = CTimer::GetCurrentTimeInCycles() / CTimer::GetCyclesPerMillisecond();
@@ -600,32 +608,40 @@ 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)
 			success = AddToLoadedVehiclesList(streamId);
+		POP_MEMID();
 		mi->RemoveRef();
 		CTxdStore::RemoveRefWithoutDelete(mi->GetTxdSlot());
 	}else{
 		// Txd
 		CTxdStore::AddRef(streamId - STREAM_OFFSET_TXD);
+		PUSH_MEMID(MEMID_STREAM_TEXUTRES);
 		success = CTxdStore::FinishLoadTxd(streamId - STREAM_OFFSET_TXD, stream);
+		POP_MEMID();
 		CTxdStore::RemoveRefWithoutDelete(streamId - STREAM_OFFSET_TXD);
 	}
 
 	RwStreamClose(stream, &mem);
-	ms_aInfoForModel[streamId].m_loadState = STREAMSTATE_LOADED;
+
+	ms_aInfoForModel[streamId].m_loadState = STREAMSTATE_LOADED;	// only done if success on PS2
+#ifndef USE_CUSTOM_ALLOCATOR
 	ms_memoryUsed += ms_aInfoForModel[streamId].GetCdSize() * CDSTREAM_SECTOR_SIZE;
+#endif
 
 	if(!success){
 		RemoveModel(streamId);
 		ReRequestModel(streamId);
-		UpdateMemoryUsed();
+		UpdateMemoryUsed();	// directly after pop on PS2
 		return false;
 	}
 
-	UpdateMemoryUsed();
+	UpdateMemoryUsed();	// directly after pop on PS2
 
 	endTime = CTimer::GetCurrentTimeInCycles() / CTimer::GetCyclesPerMillisecond();
 	timeDiff = endTime - startTime;
@@ -858,7 +874,11 @@ CStreaming::RemoveModel(int32 id)
 			CModelInfo::GetModelInfo(id)->DeleteRwObject();
 		else
 			CTxdStore::RemoveTxd(id - STREAM_OFFSET_TXD);
+#ifdef USE_CUSTOM_ALLOCATOR
+		UpdateMemoryUsed();
+#else
 		ms_memoryUsed -= ms_aInfoForModel[id].GetCdSize()*CDSTREAM_SECTOR_SIZE;
+#endif
 	}
 
 	if(ms_aInfoForModel[id].m_next){
@@ -880,6 +900,9 @@ CStreaming::RemoveModel(int32 id)
 			RpClumpGtaCancelStream();
 		else
 			CTxdStore::RemoveTxd(id - STREAM_OFFSET_TXD);
+#ifdef USE_CUSTOM_ALLOCATOR
+		UpdateMemoryUsed();
+#endif
 	}
 
 	ms_aInfoForModel[id].m_loadState = STREAMSTATE_NOTLOADED;
@@ -2044,19 +2067,25 @@ CStreaming::FlushRequestList(void)
 void
 CStreaming::ImGonnaUseStreamingMemory(void)
 {
-	// empty
+	PUSH_MEMID(MEMID_STREAM);
 }
 
 void
 CStreaming::IHaveUsedStreamingMemory(void)
 {
+	POP_MEMID();
 	UpdateMemoryUsed();
 }
 
 void
 CStreaming::UpdateMemoryUsed(void)
 {
-	// empty
+#ifdef USE_CUSTOM_ALLOCATOR
+	ms_memoryUsed =
+		gMainHeap.GetMemoryUsed(MEMID_STREAM) +
+		gMainHeap.GetMemoryUsed(MEMID_STREAM_MODELS) +
+		gMainHeap.GetMemoryUsed(MEMID_STREAM_TEXUTRES);
+#endif
 }
 
 #define STREAM_DIST 80.0f