another patch from baby.lueshi: DiscScrubber doesn't modify original file anymore, and just returns the "real" or "free'd" buffer inline with compression - so it should be a little faster too :)

git-svn-id: https://dolphin-emu.googlecode.com/svn/trunk@5032 8ced0084-cf51-0410-be5f-012b33b47a6e

3 files changed, 62 insertions, 78 deletions

diff --git a/Source/Core/DiscIO/Src/CompressedBlob.cpp b/Source/Core/DiscIO/Src/CompressedBlob.cpp
index 283633ae85..f28913c151 100644
--- a/Source/Core/DiscIO/Src/CompressedBlob.cpp
+++ b/Source/Core/DiscIO/Src/CompressedBlob.cpp
@@ -158,6 +158,8 @@ void CompressedBlobReader::GetBlock(u64 block_num, u8 *out_ptr)
 bool CompressFileToBlob(const char* infile, const char* outfile, u32 sub_type,
 						int block_size, CompressCB callback, void* arg)
 {
+	bool scrubbing = false;
+
 	if (IsCompressedBlob(infile))
 	{
 		PanicAlert("%s is already compressed! Cannot compress it further.", infile);
@@ -166,11 +168,13 @@ bool CompressFileToBlob(const char* infile, const char* outfile, u32 sub_type,
 
 	if (sub_type == 1)
 	{
-		if (!DiscScrubber::Scrub(infile, callback, arg))
+		if (!DiscScrubber::SetupScrub(infile, block_size))
 		{
 			PanicAlert("%s failed to be scrubbed. Probably the image is corrupt.", infile);
 			return false;
 		}
+
+		scrubbing = true;
 	}
 
 	FILE* inf = fopen(infile, "rb");
@@ -230,7 +234,10 @@ bool CompressFileToBlob(const char* infile, const char* outfile, u32 sub_type,
 		// u64 start = i * header.block_size;
 		// u64 size = header.block_size;
 		std::fill(in_buf, in_buf + header.block_size, 0);
-		fread(in_buf, header.block_size, 1, inf);
+		if (scrubbing)
+			DiscScrubber::GetNextBlock(inf, in_buf);
+		else
+			fread(in_buf, header.block_size, 1, inf);
 		z_stream z;
 		memset(&z, 0, sizeof(z));
 		z.zalloc = Z_NULL;
@@ -289,6 +296,7 @@ cleanup:
 	delete[] hashes;
 	fclose(f);
 	fclose(inf);
+	DiscScrubber::Cleanup();
 	callback("Done compressing disc image.", 1.0f, arg);
 	return true;
 }
diff --git a/Source/Core/DiscIO/Src/DiscScrubber.cpp b/Source/Core/DiscIO/Src/DiscScrubber.cpp
index b6674bce7d..3926b4adff 100644
--- a/Source/Core/DiscIO/Src/DiscScrubber.cpp
+++ b/Source/Core/DiscIO/Src/DiscScrubber.cpp
@@ -32,12 +32,15 @@ namespace DiscScrubber
 
 #define	CLUSTER_SIZE 0x8000
 
-static u8* m_Sector1;
-static u8* m_FreeTable;
-static u64 m_FileSize;
+u8* m_FreeTable = NULL;
+u64 m_FileSize;
+u64 m_BlockCount;
+u32 m_BlockSize;
+int m_BlocksPerCluster;
+bool m_isScrubbing = false;
 
-static std::string m_Filename;
-static IVolume* m_Disc = NULL;
+std::string m_Filename;
+IVolume* m_Disc = NULL;
 
 struct SPartitionHeader
 {
@@ -72,12 +75,11 @@ struct SPartitionGroup
 	u64 PartitionsOffset;
 	std::vector<SPartition> PartitionsVec;
 };
-static SPartitionGroup PartitionGroup[4];
+SPartitionGroup PartitionGroup[4];
 
 
 void MarkAsUsed(u64 _Offset, u64 _Size);
 void MarkAsUsedE(u64 _PartitionDataOffset, u64 _Offset, u64 _Size);
-bool MarkAsScrubbed(const char* filename);
 void ReadFromDisc(u64 _Offset, u64 _Length, u32& _Buffer);
 void ReadFromDisc(u64 _Offset, u64 _Length, u64& _Buffer);
 void ReadFromVolume(u64 _Offset, u64 _Length, u32& _Buffer);
@@ -102,10 +104,19 @@ u32 IsScrubbed(const char* filename)
 	return ScrubbedFlag;
 }
 
-bool Scrub(const char* filename, CompressCB callback, void* arg)
+bool SetupScrub(const char* filename, int block_size)
 {
 	bool success = true;
 	m_Filename = std::string(filename);
+	m_BlockSize = block_size;
+
+	if (CLUSTER_SIZE % m_BlockSize != 0)
+	{
+		ERROR_LOG(DISCIO, "block size %i is not a factor of 0x8000, scrubbing not possible", m_BlockSize);
+		return false;
+	}
+
+	m_BlocksPerCluster = CLUSTER_SIZE / m_BlockSize;
 
 	u32 version = IsScrubbed(filename);
 	if (version && version < SCRUBBER_VERSION)
@@ -115,7 +126,6 @@ bool Scrub(const char* filename, CompressCB callback, void* arg)
 	}
 	else if (version)
 	{
-		callback("DiscScrubber: This disc is already scrubbed", 0, arg);
 		NOTICE_LOG(DISCIO, "%s is already scrubbed, skipping...", filename);
 		return success;
 	}
@@ -132,73 +142,51 @@ bool Scrub(const char* filename, CompressCB callback, void* arg)
 	// Table of free blocks
 	m_FreeTable = new u8[numClusters];
 	std::fill(m_FreeTable, m_FreeTable + numClusters, 1);
-	// Fill the dummy all 1 block
-	m_Sector1 = new u8[CLUSTER_SIZE];
-	std::fill(m_Sector1, m_Sector1 + CLUSTER_SIZE, 0xFF);
 
 	// Fill out table of free blocks
-	callback("DiscScrubber: Parsing...", 0, arg);
 	success = ParseDisc();
 	// Done with it; need it closed for the next part
 	delete m_Disc;
 	m_Disc = NULL;
-
-	// Open file
-	FILE* pFile = fopen(filename, "r+b");
-	if (!pFile)
-	{
-		ERROR_LOG(DISCIO, "DiscScrubber failed to open %s", filename);
-		success = false;
-	}
+	m_BlockCount = 0;
 
 	// Let's not touch the file if we've failed up to here :p
 	if (!success)
-		goto cleanup;
+		Cleanup();
 
-	// Modify file, obeying the table of free blocks
-	NOTICE_LOG(DISCIO, "Removing garbage data...go get some coffee :)");
-	for (u32 i = 0;	i < numClusters; i++)
-	{
-		u64 CurrentOffset = (u64)i * CLUSTER_SIZE;
+	m_isScrubbing = success;
+	return success;
+}
 
-		if (m_FreeTable[i])
-		{
-			DEBUG_LOG(DISCIO, "Freeing 0x%016llx", CurrentOffset);
-			// Area is unused so fill with 1s
- 			fseek(pFile, CurrentOffset, SEEK_SET);
- 			success |= fwrite(m_Sector1, CLUSTER_SIZE, 1, pFile) == CLUSTER_SIZE;
-
-			if (!success)
-			{
-				PanicAlert("DiscScrubber failure");
-				break;
-			}
-		}
-		else
-		{
-			DEBUG_LOG(DISCIO, "Used    0x%016llx", CurrentOffset);
-		}
+void GetNextBlock(FILE* in, u8* buffer)
+{
+	u64 CurrentOffset = m_BlockCount * m_BlockSize;
+	u64 i = CurrentOffset / CLUSTER_SIZE;
 
-		// Update progress dialog
-		if (i % (numClusters / 1000) == 0)
-		{
-			char temp[512];
-			sprintf(temp, "DiscScrubber: %x/%x (%s)", i, numClusters, m_FreeTable[i] ? "Free" : "Used");
-			callback(temp, (float)i / (float)numClusters, arg);
-		}
+	if (m_isScrubbing && m_FreeTable[i])
+	{
+		DEBUG_LOG(DISCIO, "Freeing 0x%016llx", CurrentOffset);
+		std::fill(buffer, buffer + m_BlockSize, 0xFF);
+		fseek(in, m_BlockSize, SEEK_CUR);
+	}
+	else
+	{
+		DEBUG_LOG(DISCIO, "Used    0x%016llx", CurrentOffset);
+		fread(buffer, m_BlockSize, 1, in);
 	}
-	NOTICE_LOG(DISCIO, "Done removing garbage data");
-
-	if (success)
-		if (!MarkAsScrubbed(filename))
-			ERROR_LOG(DISCIO, "Really weird - failed to mark scrubbed disk as scrubbed :s");
 
-cleanup:
-	if (pFile) fclose(pFile);
-	delete[] m_Sector1;
-	delete[] m_FreeTable;
+	m_BlockCount++;
+}
 
-	return success;
+void Cleanup()
+{
+	if (m_FreeTable) delete[] m_FreeTable;
+	m_FreeTable = NULL;
+	m_FileSize = 0;
+	m_BlockCount = 0;
+	m_BlockSize = 0;
+	m_BlocksPerCluster = 0;
+	m_isScrubbing = false;
 }
 
 void MarkAsUsed(u64 _Offset, u64 _Size)
@@ -233,20 +221,6 @@ void MarkAsUsedE(u64 _PartitionDataOffset, u64 _Offset, u64 _Size)
 	MarkAsUsed(Offset, Size);
 }
 
-bool MarkAsScrubbed(const char* filename)
-{
-	FILE* f = fopen(filename, "r+b");
-
-	if (!f)
-		return false;
-
-	u8 ScrubbedFlag[1] = {SCRUBBER_VERSION};
-	fseek(f, 0x80, SEEK_SET);
-	bool success = fwrite(ScrubbedFlag, 1, 1, f) == 1;
-	fclose(f);
-	return success;
-}
-
 // Helper functions for RAW reading the BE discs
 void ReadFromDisc(u64 _Offset, u64 _Length, u32& _Buffer)
 {
diff --git a/Source/Core/DiscIO/Src/DiscScrubber.h b/Source/Core/DiscIO/Src/DiscScrubber.h
index 9180b57031..38b80b3275 100644
--- a/Source/Core/DiscIO/Src/DiscScrubber.h
+++ b/Source/Core/DiscIO/Src/DiscScrubber.h
@@ -39,7 +39,9 @@ namespace DiscScrubber
 
 u32 IsScrubbed(const char* filename);
 
-bool Scrub(const char* filename, CompressCB callback = 0, void* arg = 0);
+bool SetupScrub(const char* filename, int block_size);
+void GetNextBlock(FILE* in, u8* buffer);
+void Cleanup();
 
 } // namespace DiscScrubber