From a6c4dd26ff7a345a1fb1ae95dd8ab171659f44be Mon Sep 17 00:00:00 2001
From: Pavel Djundik <xPaw@users.noreply.github.com>
Date: Sun, 7 Jan 2024 17:35:26 +0200
Subject: [PATCH] Add `GetMemoryMappedStreamIfPossible`

Fixes #10
---
 ValvePak/ValvePak.Test/MemoryMapTest.cs | 38 +++++++++++++++++++++++++
 ValvePak/ValvePak/Package.Read.cs       | 26 +++++++++++++++++
 ValvePak/ValvePak/Package.cs            | 18 ++++++++++--
 3 files changed, 79 insertions(+), 3 deletions(-)
 create mode 100644 ValvePak/ValvePak.Test/MemoryMapTest.cs

diff --git a/ValvePak/ValvePak.Test/MemoryMapTest.cs b/ValvePak/ValvePak.Test/MemoryMapTest.cs
new file mode 100644
index 0000000..c41074e
--- /dev/null
+++ b/ValvePak/ValvePak.Test/MemoryMapTest.cs
@@ -0,0 +1,38 @@
+using System;
+using System.IO;
+using System.IO.MemoryMappedFiles;
+using NUnit.Framework;
+using SteamDatabase.ValvePak;
+
+namespace Tests
+{
+	[TestFixture]
+	public class MemoryMappedTest
+	{
+		[Test]
+		public void ReturnsMemoryMappedViewStream()
+		{
+			var path = Path.Combine(TestContext.CurrentContext.TestDirectory, "Files", "steamdb_test_dir.vpk");
+
+			using var package = new Package();
+			package.Read(path);
+
+			using var stream = package.GetMemoryMappedStreamIfPossible(package.FindEntry("kitten.jpg"));
+
+			Assert.That(stream, Is.InstanceOf<MemoryMappedViewStream>());
+		}
+
+		[Test]
+		public void ReturnsMemoryStream()
+		{
+			var path = Path.Combine(TestContext.CurrentContext.TestDirectory, "Files", "steamdb_test_single.vpk");
+
+			using var package = new Package();
+			package.Read(path);
+
+			using var stream = package.GetMemoryMappedStreamIfPossible(package.FindEntry("kitten.jpg"));
+
+			Assert.That(stream, Is.InstanceOf<MemoryStream>());
+		}
+	}
+}
diff --git a/ValvePak/ValvePak/Package.Read.cs b/ValvePak/ValvePak/Package.Read.cs
index 59862eb..f05a5dc 100644
--- a/ValvePak/ValvePak/Package.Read.cs
+++ b/ValvePak/ValvePak/Package.Read.cs
@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.IO;
+using System.IO.MemoryMappedFiles;
 using System.Runtime.CompilerServices;
 using System.Text;
 
@@ -343,6 +344,31 @@ private Stream GetFileStream(ushort archiveIndex)
 			return stream;
 		}
 
+		/// <summary>
+		/// Returns <see cref="MemoryMappedViewStream"/> when possible, otherwise reads entry into a byte array and returns <see cref="MemoryStream"/>.
+		/// This only works on split packages (<see cref="IsDirVPK"/>) and when entries have no preload bytes.
+		/// </summary>
+		/// <param name="entry">Package entry.</param>
+		/// <returns>Stream for a given package entry contents.</returns>
+		public Stream GetMemoryMappedStreamIfPossible(PackageEntry entry)
+		{
+			if (!IsDirVPK || entry.ArchiveIndex == 0x7FFF || entry.SmallData.Length > 0)
+			{
+				ReadEntry(entry, out var output, false);
+
+				return new MemoryStream(output);
+			}
+
+			if (!MemoryMappedPaks.TryGetValue(entry.ArchiveIndex, out var stream))
+			{
+				var path = $"{FileName}_{entry.ArchiveIndex:D3}.vpk";
+				stream = MemoryMappedFile.CreateFromFile(path, FileMode.Open, null, 0, MemoryMappedFileAccess.Read);
+				MemoryMappedPaks[entry.ArchiveIndex] = stream;
+			}
+
+			return stream.CreateViewStream(entry.Offset, entry.Length, MemoryMappedFileAccess.Read);
+		}
+
 		[MethodImpl(MethodImplOptions.AggressiveInlining)]
 		private string ReadNullTermUtf8String(MemoryStream ms)
 		{
diff --git a/ValvePak/ValvePak/Package.cs b/ValvePak/ValvePak/Package.cs
index e2108de..ef10cf5 100644
--- a/ValvePak/ValvePak/Package.cs
+++ b/ValvePak/ValvePak/Package.cs
@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.IO;
+using System.IO.MemoryMappedFiles;
 
 namespace SteamDatabase.ValvePak
 {
@@ -17,6 +18,7 @@ public partial class Package : IDisposable
 		public const char DirectorySeparatorChar = '/';
 
 		private BinaryReader Reader;
+		private readonly Dictionary<int, MemoryMappedFile> MemoryMappedPaks = [];
 
 		/// <summary>
 		/// Gets the file name.
@@ -111,10 +113,20 @@ public void Dispose()
 
 		protected virtual void Dispose(bool disposing)
 		{
-			if (disposing && Reader != null)
+			if (disposing)
 			{
-				Reader.Dispose();
-				Reader = null;
+				if (Reader != null)
+				{
+					Reader.Dispose();
+					Reader = null;
+				}
+
+				foreach (var stream in MemoryMappedPaks.Values)
+				{
+					stream.Dispose();
+				}
+
+				MemoryMappedPaks.Clear();
 			}
 		}