p4k extract speedup

src/StarBreaker.Cli/packages.lock.json

@@ -8,6 +8,12 @@
         "resolved": "2.3.5",
         "contentHash": "WaZDt7FC1ViEzS5m7w7lOWT/aluZ28fFkPcRdkQ+7DowqjTmr5YZswNIDqSuBDaeGXPe8VzTZKMltKrXhIOt9Q=="
       },
+      "Microsoft.DotNet.ILCompiler": {
+        "type": "Direct",
+        "requested": "[9.0.0, )",
+        "resolved": "9.0.0",
+        "contentHash": "bbnlV2PbUmCQ8Ndpx0kJaicLyV28IU+4IzyctQLL57+DxrHurYr2qsJrC8+yD44Q0DyPfv2oM168c1Tk6Bxbmg=="
+      },
       "Microsoft.NET.ILLink.Tasks": {
         "type": "Direct",
         "requested": "[9.0.0, )",

src/StarBreaker.Common/StreamExtensions.cs

@@ -23,42 +23,16 @@ public static T Read<T>(this Stream stream) where T : unmanaged
 
     public static void CopyAmountTo(this Stream source, Stream destination, int byteCount)
     {
-        var buffer = ArrayPool<byte>.Shared.Rent(byteCount);
+        var rent = ArrayPool<byte>.Shared.Rent(byteCount);
+        var buffer = rent.AsSpan(0, byteCount);
         try
         {
-            while (byteCount > 0)
-            {
-                var n = source.Read(buffer, 0, Math.Min(byteCount, buffer.Length));
-                if (n == 0)
-                    throw new Exception("Failed to read from stream");
-                destination.Write(buffer, 0, n);
-                byteCount -= n;
-            }
+            source.ReadExactly(buffer);
+            destination.Write(buffer);
         }
         finally
         {
-            ArrayPool<byte>.Shared.Return(buffer);
-        }
-    }
-
-    public static byte[] ToArray(this Stream stream)
-    {
-        if (stream is MemoryStream ms)
-            return ms.ToArray();
-
-        try
-        {
-            var count = stream.Position;
-            var buffer = new byte[count];
-            stream.Position = 0;
-            stream.ReadExactly(buffer, 0, buffer.Length);
-            return buffer;
-        }
-        catch (NotSupportedException)
-        {
-            using var m = new MemoryStream();
-            stream.CopyTo(m);
-            return m.ToArray();
+            ArrayPool<byte>.Shared.Return(rent);
         }
     }
 }

src/StarBreaker.P4k/P4kExtractor.cs

@@ -27,7 +27,7 @@ public void Extract(string outputDir, string? filter = null, IProgress<double>?
     {
         //TODO: if the filter is for *.dds, make sure to include *.dds.N too. Maybe do the pre processing before we filter?
         var filteredEntries = (filter is null
-            ? _p4KFile.Entries.OrderByDescending(entry => entry.UncompressedSize).Take(1000).ToArray()
+            ? _p4KFile.Entries.ToArray()
             : _p4KFile.Entries.Where(entry => FileSystemName.MatchesSimpleExpression(filter, entry.Name))).ToArray();
 
         var numberOfEntries = filteredEntries.Length;
@@ -46,10 +46,6 @@ public void Extract(string outputDir, string? filter = null, IProgress<double>?
         // 5. find multiple file -> single file unsplit processors - remove from the list so we don't double process
         // run it!
         Parallel.ForEach(filteredEntries,
-            new ParallelOptions
-            {
-                MaxDegreeOfParallelism = 1
-            },
             entry =>
             {
                 if (entry.UncompressedSize == 0)
@@ -63,7 +59,7 @@ public void Extract(string outputDir, string? filter = null, IProgress<double>?
                 using (var writeStream = new FileStream(entryPath, FileMode.Create, FileAccess.Write, FileShare.None,
                            bufferSize: entry.UncompressedSize > int.MaxValue ? 81920 : (int)entry.UncompressedSize, useAsync: true))
                 {
-                    using (var entryStream = _p4KFile.Open(entry))
+                    using (var entryStream = _p4KFile.OpenStream(entry))
                     {
                         entryStream.CopyTo(writeStream);
                     }
@@ -77,6 +73,9 @@ public void Extract(string outputDir, string? filter = null, IProgress<double>?
                         progress?.Report(processedEntries / (double)numberOfEntries);
                     }
                 }
-            });
+            }
+        );
+
+        progress?.Report(1);
     }
 }

src/StarBreaker.P4k/P4kFile.FileSystem.cs

@@ -87,6 +87,7 @@ public Stream OpenRead(string path)
         if (current.ZipEntry == null)
             throw new FileNotFoundException();
 
-        return Open(current.ZipEntry);
+        //Is this a bad idea? Most things that use this rely on the stream being seekable.
+        return new MemoryStream(OpenInMemory(current.ZipEntry));
     }
 }

src/StarBreaker.P4k/P4kFile.cs

@@ -10,6 +10,7 @@ namespace StarBreaker.P4k;
 
 public sealed partial class P4kFile
 {
+    [ThreadStatic] private static Decompressor? decompressor;
     private readonly Aes _aes;
 
     public ZipEntry[] Entries { get; }
@@ -176,7 +177,8 @@ private static ZipEntry ReadEntry(BinaryReader reader)
         }
     }
 
-    public Stream Open(ZipEntry entry)
+    // Represents the raw stream from the p4k file, before any decryption or decompression
+    private StreamSegment OpenInternal(ZipEntry entry)
     {
         var p4kStream = new FileStream(P4KPath, FileMode.Open, FileAccess.Read, FileShare.Read);
 
@@ -187,88 +189,64 @@ public Stream Open(ZipEntry entry)
         var localHeader = p4kStream.Read<LocalFileHeader>();
 
         p4kStream.Seek(localHeader.FileNameLength + localHeader.ExtraFieldLength, SeekOrigin.Current);
-        Stream entryStream = new StreamSegment(p4kStream, p4kStream.Position, (long)entry.CompressedSize, false);
+
+        return new StreamSegment(p4kStream, p4kStream.Position, (long)entry.CompressedSize, false);
+    }
+
+    // Remarks: Streams returned by this method might not support seeking or length.
+    // If these are required, consider using OpenInMemory instead.
+    public Stream OpenStream(ZipEntry entry)
+    {
+        // Represents the raw stream from the p4k file, before any decryption or decompression
+        var entryStream = OpenInternal(entry);
 
         return entry switch
         {
-            { IsCrypted: true, CompressionMethod: 100 } => GetDecryptStream(entryStream, entry.UncompressedSize),
-            { IsCrypted: false, CompressionMethod: 100 } => GetDecompressionStream(entryStream, entry.UncompressedSize),
+            { IsCrypted: true, CompressionMethod: 100 } => GetDecryptStream(entryStream, entry.CompressedSize),
+            { IsCrypted: false, CompressionMethod: 100 } => GetDecompressionStream(entryStream),
             { IsCrypted: false, CompressionMethod: 0 } when entry.CompressedSize != entry.UncompressedSize => throw new Exception("Invalid stored file"),
             { IsCrypted: false, CompressionMethod: 0 } => entryStream,
             _ => throw new Exception("Invalid compression method")
         };
     }
 
-    private MemoryStream GetDecryptStream(Stream entryStream, ulong uncompressedSize)
+    private DecompressionStream GetDecryptStream(Stream entryStream, ulong compressedSize)
     {
         using var transform = _aes.CreateDecryptor();
+        var ms = new MemoryStream((int)compressedSize);
+        using (var crypto = new CryptoStream(entryStream, transform, CryptoStreamMode.Read))
+            crypto.CopyTo(ms);
 
-        var rent = ArrayPool<byte>.Shared.Rent((int)entryStream.Length);
-        try
-        {
-            using var rented = new MemoryStream(rent, 0, (int)entryStream.Length, true, true);
-            using (var crypto = new CryptoStream(entryStream, transform, CryptoStreamMode.Read))
-            {
-                crypto.CopyTo(rented);
-            }
-
-            // Trim NULL off end of stream
-            rented.Seek(-1, SeekOrigin.End);
-            while (rented.Position > 1 && rented.ReadByte() == 0)
-                rented.Seek(-2, SeekOrigin.Current);
-            rented.SetLength(rented.Position);
+        // Trim NULL off end of stream
+        ms.Seek(-1, SeekOrigin.End);
+        while (ms.Position > 1 && ms.ReadByte() == 0)
+            ms.Seek(-2, SeekOrigin.Current);
+        ms.SetLength(ms.Position);
 
-            rented.Seek(0, SeekOrigin.Begin);
+        ms.Seek(0, SeekOrigin.Begin);
 
-            using var decompressionStream = new DecompressionStream(rented, leaveOpen: true);
-            var finalStream = new MemoryStream((int)uncompressedSize);
-            decompressionStream.CopyTo(finalStream);
-
-            finalStream.Seek(0, SeekOrigin.Begin);
-
-            return finalStream;
-        }
-        finally
-        {
-            ArrayPool<byte>.Shared.Return(rent);
-        }
+        return GetDecompressionStream(ms);
     }
 
-    private static MemoryStream GetDecompressionStream(Stream entryStream, ulong uncompressedSize)
+    private static DecompressionStream GetDecompressionStream(Stream entryStream)
     {
-        // We need to make a new memoryStream and copy the data over.
-        // This is because the decompression stream doesn't support seeking/position/length.
-
-        var buffer = new MemoryStream((int)uncompressedSize);
-
-        //close the entryStream (p4k file probably) when we're done with it
-        using (var decompressionStream = new DecompressionStream(entryStream, leaveOpen: false))
-            decompressionStream.CopyTo(buffer);
-
-        buffer.Seek(0, SeekOrigin.Begin);
-
-        return buffer;
+        return new DecompressionStream(entryStream, decompressor: decompressor ??= new Decompressor());
     }
 
-    public static List<ZipExtraField> ReadExtraFields(BinaryReader br, ushort length)
+    public byte[] OpenInMemory(ZipEntry entry)
     {
-        var fields = new List<ZipExtraField>();
+        if (entry.UncompressedSize > int.MaxValue)
+            throw new Exception("File too large to load into memory. Use OpenStream instead");
 
-        while (length > 0)
-        {
-            var tag = br.ReadUInt16();
-            var size = br.ReadUInt16();
-            var data = br.ReadBytes(size - 4);
+        var uncompressedSize = checked((int)entry.UncompressedSize);
 
-            fields.Add(new ZipExtraField(tag, size, data));
-            length -= size;
-        }
+        var ms = new MemoryStream(uncompressedSize);
+        OpenStream(entry).CopyTo(ms);
 
-        return fields;
-    }
+        // If the stream is larger than the uncompressed size, trim it.
+        // This can happen because of decryption padding bytes :(
+        ms.SetLength(ms.Position);
 
-    public int GetBufferSize(ulong size)
-    {
-        return Math.Min(81920, size > int.MaxValue ? 81920 : (int)size);
+        return ms.ToArray();
     }
 }

src/StarBreaker/Screens/Tabs/P4kTabView/P4kTabViewModel.cs

@@ -85,48 +85,53 @@ private void SelectionChanged(object? sender, TreeSelectionModelSelectionChanged
             return;
         }
 
-        //if above 1gb, don't load
-        if (selectedEntry.ZipEntry.UncompressedSize > 1024 * 1024 * 1024)
+        if (selectedEntry.ZipEntry.UncompressedSize > int.MaxValue)
+        {
+            _logger.LogWarning("File too big to preview");
             return;
+        }
 
         //todo: for a big ass file show a loading screen or something
         Preview = null;
         Task.Run(() =>
         {
-            //TODO: move this to a service?
-            byte[] buffer;
-            using (var stream = _p4KService.P4kFile.Open(selectedEntry.ZipEntry))
-                buffer = stream.ToArray();
+            try
+            {
+                //TODO: move this to a service?
+                var buffer = _p4KService.P4kFile.OpenInMemory(selectedEntry.ZipEntry);
 
-            FilePreviewViewModel preview;
+                FilePreviewViewModel preview;
 
-            //check cryxml before extension since ".xml" sometimes is cxml sometimes plaintext
-            if (CryXmlB.CryXml.TryOpen(new MemoryStream(buffer), out var c))
-            {
-                _logger.LogInformation("cryxml");
-                var stringwriter = new StringWriter();
-                c.WriteTo(XmlWriter.Create(stringwriter, new XmlWriterSettings{Indent = true}));
-                preview = new TextPreviewViewModel(stringwriter.ToString());
-            }
-            else if (plaintextExtensions.Any(p => selectedEntry.GetName().EndsWith(p, StringComparison.InvariantCultureIgnoreCase)))
-            {
-                _logger.LogInformation("plaintextExtensions");
+                //check cryxml before extension since ".xml" sometimes is cxml sometimes plaintext
+                if (CryXmlB.CryXml.TryOpen(new MemoryStream(buffer), out var c))
+                {
+                    _logger.LogInformation("cryxml");
+                    preview = new TextPreviewViewModel(c.ToString());
+                }
+                else if (plaintextExtensions.Any(p => selectedEntry.GetName().EndsWith(p, StringComparison.InvariantCultureIgnoreCase)))
+                {
+                    _logger.LogInformation("plaintextExtensions");
 
-                preview = new TextPreviewViewModel(Encoding.UTF8.GetString(buffer));
-            }
-            else if (ddsLodExtensions.Any(p => selectedEntry.GetName().EndsWith(p, StringComparison.InvariantCultureIgnoreCase)))
-            {
-                _logger.LogInformation("ddsLodExtensions");
-                preview = new DdsPreviewViewModel(buffer);
+                    preview = new TextPreviewViewModel(Encoding.UTF8.GetString(buffer));
+                }
+                else if (ddsLodExtensions.Any(p => selectedEntry.GetName().EndsWith(p, StringComparison.InvariantCultureIgnoreCase)))
+                {
+                    _logger.LogInformation("ddsLodExtensions");
+                    preview = new DdsPreviewViewModel(buffer);
+                }
+                else
+                {
+                    _logger.LogInformation("hex");
+                    preview = new HexPreviewViewModel(buffer);
+                }
+                //todo other types
+
+                Dispatcher.UIThread.Post(() => Preview = preview);
             }
-            else
+            catch (Exception exception)
             {
-                _logger.LogInformation("hex");
-                preview = new HexPreviewViewModel(buffer);
+                _logger.LogError(exception, "Failed to preview file");
             }
-            //todo other types
-
-            Dispatcher.UIThread.Post(() => Preview = preview);
         });
     }
 }

src/StarBreaker/Screens/Tabs/P4kTabView/Previews/HexPreview/HexPreviewViewModel.cs

@@ -7,6 +7,7 @@ public  sealed partial class HexPreviewViewModel : FilePreviewViewModel
 {
     public HexPreviewViewModel(byte[] data)
     {
+        //TODO: make this use a Stream instead of a byte array?
         Document = new MemoryBinaryDocument(data, true);
     }