From 854bdbcaf3f22b414ee54b5f4ff228f6c4b43b17 Mon Sep 17 00:00:00 2001 From: Timothy Baldridge Date: Wed, 4 Mar 2020 21:32:38 -0700 Subject: [PATCH] Convert remaining BSA types to us virtual memory --- Compression.BSA/BA2Builder.cs | 4 +-- Compression.BSA/BSABuilder.cs | 64 ++++++++++++++++++++-------------- Compression.BSA/TES3Builder.cs | 9 +++-- Compression.BSA/Utils.cs | 2 ++ 4 files changed, 46 insertions(+), 33 deletions(-) diff --git a/Compression.BSA/BA2Builder.cs b/Compression.BSA/BA2Builder.cs index 27d2ab8f..813a1163 100644 --- a/Compression.BSA/BA2Builder.cs +++ b/Compression.BSA/BA2Builder.cs @@ -203,7 +203,7 @@ namespace Compression.BSA bw.BaseStream.Position = _offsetOffset; bw.Write((ulong)pos); bw.BaseStream.Position = pos; - _dataSlab.CopyTo(bw.BaseStream); + _dataSlab.CopyToLimit(bw.BaseStream, (int)_dataSlab.Length); } } @@ -267,7 +267,7 @@ namespace Compression.BSA wtr.Write((ulong)pos); wtr.BaseStream.Position = pos; _dataSrc.Position = 0; - _dataSrc.CopyTo(wtr.BaseStream); + _dataSrc.CopyToLimit(wtr.BaseStream, (int)_dataSrc.Length); } } } diff --git a/Compression.BSA/BSABuilder.cs b/Compression.BSA/BSABuilder.cs index 375f530c..6569be09 100644 --- a/Compression.BSA/BSABuilder.cs +++ b/Compression.BSA/BSABuilder.cs @@ -6,6 +6,7 @@ using System.Text; using ICSharpCode.SharpZipLib.Zip.Compression.Streams; using K4os.Compression.LZ4; using K4os.Compression.LZ4.Streams; +using Wabbajack.Common; using File = Alphaleonis.Win32.Filesystem.File; using Path = Alphaleonis.Win32.Filesystem.Path; @@ -22,11 +23,13 @@ namespace Compression.BSA internal uint _offset; internal uint _totalFileNameLength; internal uint _version; + internal DiskSlabAllocator _slab; public BSABuilder() { _fileId = Encoding.ASCII.GetBytes("BSA\0"); _offset = 0x24; + _slab = new DiskSlabAllocator(); } public BSABuilder(BSAStateObject bsaStateObject) : this() @@ -74,6 +77,7 @@ namespace Compression.BSA public void Dispose() { + _slab.Dispose(); } public void AddFile(FileStateObject state, Stream src) { @@ -231,7 +235,7 @@ namespace Compression.BSA internal string _path; private byte[] _pathBSBytes; internal byte[] _pathBytes; - internal byte[] _rawData; + private Stream _srcData; public static FileEntry Create(BSABuilder bsa, string path, Stream src, bool flipCompression) { @@ -244,11 +248,9 @@ namespace Compression.BSA entry._pathBytes = entry._path.ToTermString(bsa.HeaderType); entry._pathBSBytes = entry._path.ToBSString(); entry._flipCompression = flipCompression; + entry._srcData = src; - var ms = new MemoryStream(); - src.CopyTo(ms); - entry._rawData = ms.ToArray(); - entry._originalSize = entry._rawData.Length; + entry._originalSize = (int)entry._srcData.Length; if (entry.Compressed) entry.CompressData(); @@ -275,29 +277,37 @@ namespace Compression.BSA private void CompressData() { - if (_bsa.HeaderType == VersionType.SSE) + switch (_bsa.HeaderType) { - var r = new MemoryStream(); - using (var w = LZ4Stream.Encode(r, new LZ4EncoderSettings {CompressionLevel = LZ4Level.L10_OPT})) + case VersionType.SSE: { - new MemoryStream(_rawData).CopyTo(w); + var r = new MemoryStream(); + using (var w = LZ4Stream.Encode(r, new LZ4EncoderSettings {CompressionLevel = LZ4Level.L10_OPT}, true)) + { + _srcData.CopyTo(w); + } + _srcData = _bsa._slab.Allocate(r.Length); + r.Position = 0; + r.CopyTo(_srcData); + _srcData.Position = 0; + break; } - - _rawData = r.ToArray(); - } - else if (_bsa.HeaderType == VersionType.FO3 || _bsa.HeaderType == VersionType.TES4) - { - var r = new MemoryStream(); - using (var w = new DeflaterOutputStream(r)) + case VersionType.FO3: + case VersionType.TES4: { - new MemoryStream(_rawData).CopyTo(w); + var r = new MemoryStream(); + using (var w = new DeflaterOutputStream(r)) + { + _srcData.CopyTo(w); + } + _srcData = _bsa._slab.Allocate(r.Length); + r.Position = 0; + r.CopyTo(_srcData); + _srcData.Position = 0; + break; } - - _rawData = r.ToArray(); - } - else - { - throw new NotImplementedException($"Can't compress data for {_bsa.HeaderType} BSAs."); + default: + throw new NotImplementedException($"Can't compress data for {_bsa.HeaderType} BSAs."); } } @@ -305,7 +315,7 @@ namespace Compression.BSA { wtr.Write(_hash); - var size = _rawData.Length; + var size = _srcData.Length; if (_bsa.HasNameBlobs) size += _pathBSBytes.Length; if (Compressed) size += 4; if (_flipCompression) @@ -329,11 +339,13 @@ namespace Compression.BSA if (Compressed) { wtr.Write((uint) _originalSize); - wtr.BaseStream.Write(_rawData, 0, _rawData.Length); + _srcData.Position = 0; + _srcData.CopyToLimit(wtr.BaseStream, (int)_srcData.Length); } else { - wtr.BaseStream.Write(_rawData, 0, _rawData.Length); + _srcData.Position = 0; + _srcData.CopyToLimit(wtr.BaseStream, (int)_srcData.Length); } } } diff --git a/Compression.BSA/TES3Builder.cs b/Compression.BSA/TES3Builder.cs index 1bf7b8c3..2fccc6a1 100644 --- a/Compression.BSA/TES3Builder.cs +++ b/Compression.BSA/TES3Builder.cs @@ -8,19 +8,18 @@ namespace Compression.BSA public class TES3Builder : IBSABuilder { private TES3ArchiveState _state; - private (TES3FileState state, byte[] data)[] _files; + private (TES3FileState state, Stream data)[] _files; public TES3Builder(TES3ArchiveState state) { _state = state; - _files = new (TES3FileState state, byte[] data)[_state.FileCount]; + _files = new (TES3FileState state, Stream data)[_state.FileCount]; } public void AddFile(FileStateObject state, Stream src) { - using var br = new BinaryReader(src); var cstate = (TES3FileState)state; - _files[state.Index] = (cstate, br.ReadBytes((int)cstate.Size)); + _files[state.Index] = (cstate, src); } public void Build(string filename) @@ -66,7 +65,7 @@ namespace Compression.BSA foreach (var (state, data) in _files) { bw.BaseStream.Position = _state.DataOffset + state.Offset; - bw.Write(data); + data.CopyTo(bw.BaseStream); } } diff --git a/Compression.BSA/Utils.cs b/Compression.BSA/Utils.cs index 64026327..9fd36c77 100644 --- a/Compression.BSA/Utils.cs +++ b/Compression.BSA/Utils.cs @@ -156,6 +156,8 @@ namespace Compression.BSA { var to_read = Math.Min(buff.Length, limit); var read = frm.Read(buff, 0, to_read); + if (read == 0) + throw new Exception("End of stream before end of limit"); tw.Write(buff, 0, read); limit -= read; }