Skip to content

Commit dd8da93

Browse files
AXiX-officialAXiX-official
committed
Optimize BlockStream Caching
1 parent 3c268fc commit dd8da93

7 files changed

Lines changed: 131 additions & 49 deletions

File tree

‎UnityAsset.NET/Asset.cs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using UnityAsset.NET.Files.SerializedFiles;
22
using UnityAsset.NET.IO.Reader;
3+
using UnityAsset.NET.IO.Stream;
34
using UnityAsset.NET.TypeTree;
45
using UnityAsset.NET.TypeTree.PreDefined;
56

@@ -38,6 +39,7 @@ public IUnityAsset Value
3839
{
3940
using var reader = DataReader;
4041
_value = UnityObjectFactory.Create(Info.Type, reader);
42+
BlockStream.OnAssetParsed(this);
4143
if (IsNamedAsset)
4244
{
4345
_name = ((INamedObject)_value).m_Name;

‎UnityAsset.NET/AssetManager.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -244,5 +244,6 @@ public void Clear()
244244
TpkUnityTreeNodeFactory.Deinit();
245245
AssemblyManager.CleanCache();
246246
BlockStream.Cache.Reset(Setting.DefaultBlockCacheSize);
247+
BlockStream.AssetToBlockCache = new();
247248
}
248249
}

‎UnityAsset.NET/Files/BundleFiles/FileEntry.cs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,12 @@ namespace UnityAsset.NET.Files.BundleFiles;
55

66
public sealed class FileEntry
77
{
8-
public Int64 Offset;
9-
public Int64 Size;
8+
public ulong Offset;
9+
public ulong Size;
1010
public UInt32 Flags;
1111
public string Path;
1212

13-
public FileEntry(Int64 offset, Int64 size, UInt32 flags, string path)
13+
public FileEntry(ulong offset, ulong size, UInt32 flags, string path)
1414
{
1515
Offset = offset;
1616
Size = size;
@@ -19,8 +19,8 @@ public FileEntry(Int64 offset, Int64 size, UInt32 flags, string path)
1919
}
2020

2121
public static FileEntry Parse(IReader reader) => new (
22-
reader.ReadInt64(),
23-
reader.ReadInt64(),
22+
reader.ReadUInt64(),
23+
reader.ReadUInt64(),
2424
reader.ReadUInt32(),
2525
reader.ReadNullTerminatedString()
2626
);

‎UnityAsset.NET/Files/SerializedFiles/SerializedFile.cs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
using System.Text;
2+
using SharpCompress.Common;
23
using UnityAsset.NET.Extensions;
34
using UnityAsset.NET.FileSystem;
45
using UnityAsset.NET.IO;
56
using UnityAsset.NET.IO.Reader;
7+
using UnityAsset.NET.IO.Stream;
68
using UnityAsset.NET.TypeTree.PreDefined;
79
using UnityAsset.NET.TypeTree.PreDefined.Types;
810

@@ -68,6 +70,16 @@ public static SerializedFile Parse(BundleFile bf, IReaderProvider readerProvider
6870
assets.Add(new Asset(sf, assetInfo));
6971
}
7072
sf.BuildMap();
73+
74+
if (readerProvider is CustomStreamReaderProvider csrp)
75+
{
76+
if (csrp.StreamProvider is FileEntryStreamProvider fesp)
77+
{
78+
if (fesp.StreamProvider is BlockStreamProvider bsp)
79+
BlockStream.RegisterAssetToBlockMap(fesp, bsp, sf);
80+
}
81+
}
82+
7183
return sf;
7284
}
7385

‎UnityAsset.NET/IO/Reader/CustomStreamReader.cs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -292,12 +292,12 @@ public void Dispose()
292292

293293
public class CustomStreamReaderProvider : IReaderProvider
294294
{
295-
private readonly IStreamProvider _streamProvider;
295+
public readonly IStreamProvider StreamProvider;
296296

297297
public CustomStreamReaderProvider(IStreamProvider streamProvider)
298298
{
299-
_streamProvider = streamProvider;
299+
StreamProvider = streamProvider;
300300
}
301301

302-
public IReader CreateReader(Endianness endian = Endianness.BigEndian) => new CustomStreamReader(_streamProvider, endian);
302+
public IReader CreateReader(Endianness endian = Endianness.BigEndian) => new CustomStreamReader(StreamProvider, endian);
303303
}

‎UnityAsset.NET/IO/Stream/BlockStream.cs‎

Lines changed: 99 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,29 @@
1-
using UnityAsset.NET.Enums;
1+
using System.Collections.Concurrent;
2+
using UnityAsset.NET.Enums;
23
using UnityAsset.NET.Files.BundleFiles;
4+
using UnityAsset.NET.Files.SerializedFiles;
5+
using UnityAsset.NET.IO.Reader;
36

47
namespace UnityAsset.NET.IO.Stream;
58

69
public class BlockStreamProvider : IStreamProvider
710
{
8-
private readonly BlockStream.BlockInfo[] _blocks;
9-
private readonly IReaderProvider _baseReaderProvider;
11+
public readonly BlockStream.BlockInfo[] Blocks;
12+
public readonly IReaderProvider BaseReaderProvider;
1013
private readonly ulong _length;
1114
private readonly UnityCN? _unityCnInfo;
1215

1316
public BlockStreamProvider(List<StorageBlockInfo> blocks, IReaderProvider readerProvider,
1417
UnityCN? unityCnInfo = null)
1518
{
16-
_blocks = new BlockStream.BlockInfo[blocks.Count];
19+
Blocks = new BlockStream.BlockInfo[blocks.Count];
1720
ulong currentOffset = 0;
1821
ulong uncompressedOffset = 0;
1922
for (int i = 0; i < blocks.Count; i++)
2023
{
2124
var block = blocks[i];
2225
var compressionType = (CompressionType)(block.Flags & StorageBlockFlags.CompressionTypeMask);
23-
_blocks[i] = new BlockStream.BlockInfo(
26+
Blocks[i] = new BlockStream.BlockInfo(
2427
block.UncompressedSize,
2528
block.CompressedSize,
2629
currentOffset,
@@ -31,12 +34,12 @@ public BlockStreamProvider(List<StorageBlockInfo> blocks, IReaderProvider reader
3134
uncompressedOffset += block.UncompressedSize;
3235
}
3336

34-
_baseReaderProvider = readerProvider;
37+
BaseReaderProvider = readerProvider;
3538
_length = uncompressedOffset;
3639
_unityCnInfo = unityCnInfo;
3740
}
3841

39-
public System.IO.Stream OpenStream() => new BlockStream(_blocks, _baseReaderProvider, _length, _unityCnInfo);
42+
public System.IO.Stream OpenStream() => new BlockStream(Blocks, BaseReaderProvider, _length, _unityCnInfo);
4043
}
4144

4245
public class BlockStream : System.IO.Stream
@@ -60,6 +63,59 @@ public readonly record struct BlockInfo(
6063
public readonly record struct CacheKey(IReaderProvider Provider, int BlockIndex);
6164

6265
public static CustomMemoryCache<CacheKey, byte[]> Cache = new (maxSize: Setting.DefaultBlockCacheSize);
66+
67+
public static ConcurrentDictionary<CacheKey, (int parsed, int total)> AssetToBlockCache = new();
68+
69+
public static void RegisterAssetToBlockMap(FileEntryStreamProvider fesp, BlockStreamProvider bsp, SerializedFile sf)
70+
{
71+
var offset = fesp.FileEntry.Offset;
72+
var blocks = bsp.Blocks;
73+
foreach (var info in sf.Metadata.AssetInfos)
74+
{
75+
var pos = offset + sf.Header.DataOffset + info.ByteOffset;
76+
var index = FindBlockIndex(blocks, pos);
77+
while (index < blocks.Length && pos + info.ByteSize > blocks[index].UncompressedOffset)
78+
{
79+
var key = new CacheKey(bsp.BaseReaderProvider, index);
80+
AssetToBlockCache.AddOrUpdate(key,
81+
addValue: (0, 1),
82+
updateValueFactory: (k, existing) =>
83+
(existing.parsed, existing.total + 1));
84+
index++;
85+
}
86+
}
87+
}
88+
89+
public static void OnAssetParsed(Asset asset)
90+
{
91+
var sf = asset.SourceFile;
92+
if (sf.ReaderProvider is CustomStreamReaderProvider csrp)
93+
{
94+
if (csrp.StreamProvider is FileEntryStreamProvider fesp)
95+
{
96+
if (fesp.StreamProvider is BlockStreamProvider bsp)
97+
{
98+
var offset = fesp.FileEntry.Offset;
99+
var blocks = bsp.Blocks;
100+
var pos = offset + sf.Header.DataOffset + asset.Info.ByteOffset;
101+
var index = FindBlockIndex(blocks, pos);
102+
while (index < blocks.Length && pos + asset.Info.ByteSize > blocks[index].UncompressedOffset)
103+
{
104+
var key = new CacheKey(bsp.BaseReaderProvider, index);
105+
var newStats = AssetToBlockCache.AddOrUpdate(key,
106+
addValue: (0, 1),
107+
updateValueFactory: (_, existing) =>
108+
(existing.parsed + 1, existing.total));
109+
if (newStats.parsed == newStats.total)
110+
{
111+
Cache.Remove(key);
112+
}
113+
index++;
114+
}
115+
}
116+
}
117+
}
118+
}
63119

64120
public BlockStream(BlockInfo[] blocks, IReaderProvider baseReaderProvider, ulong length, UnityCN? unityCnInfo = null)
65121
{
@@ -80,18 +136,16 @@ public override long Position
80136
set => Seek(value, SeekOrigin.Begin);
81137
}
82138

83-
private void EnsureBlockLoaded(ulong position)
139+
private static int FindBlockIndex(BlockInfo[] blocks, ulong position)
84140
{
85-
if (_disposed) throw new ObjectDisposedException(nameof(BlockStream));
86-
87141
int low = 0;
88-
int high = Blocks.Length - 1;
142+
int high = blocks.Length - 1;
89143
int blockIndex = -1;
90-
144+
91145
while (low <= high)
92146
{
93147
int mid = low + (high - low) / 2;
94-
var block = Blocks[mid];
148+
var block = blocks[mid];
95149
if (position >= block.UncompressedOffset && position < block.UncompressedOffset + block.UncompressedSize)
96150
{
97151
blockIndex = mid;
@@ -107,6 +161,30 @@ private void EnsureBlockLoaded(ulong position)
107161
}
108162
}
109163

164+
return blockIndex;
165+
}
166+
167+
private byte[] DecompressBlock(int index)
168+
{
169+
var block = Blocks[index];
170+
using var reader = _baseReaderProvider.CreateReader();
171+
reader.Position = (long)block.CompressedOffset;
172+
var compressedData = reader.ReadBytes((int)block.CompressedSize);
173+
if (_unityCnInfo != null && (block.CompressionType == CompressionType.Lz4 || block.CompressionType == CompressionType.Lz4HC))
174+
{
175+
_unityCnInfo.DecryptBlock(compressedData, compressedData.Length, index);
176+
}
177+
178+
var uncompressedBuffer = new byte[block.UncompressedSize];
179+
Compression.DecompressToBytes(compressedData, uncompressedBuffer.AsSpan(0, (int)block.UncompressedSize), block.CompressionType);
180+
181+
return uncompressedBuffer;
182+
}
183+
184+
private void EnsureBlockLoaded(ulong position)
185+
{
186+
var blockIndex = FindBlockIndex(Blocks, position);
187+
110188
if (blockIndex == -1)
111189
{
112190
throw new ArgumentOutOfRangeException(nameof(position), "Position is beyond the end of the data");
@@ -115,26 +193,15 @@ private void EnsureBlockLoaded(ulong position)
115193
if (blockIndex != CurrentBlockIndex)
116194
{
117195
var block = Blocks[blockIndex];
196+
var key = new CacheKey(_baseReaderProvider, blockIndex);
118197

119-
var blockBuffer = Cache.GetOrCreate(
120-
key: new(_baseReaderProvider, blockIndex),
121-
factory: () =>
122-
{
123-
using var reader = _baseReaderProvider.CreateReader();
124-
reader.Position = (long)block.CompressedOffset;
125-
var compressedData = reader.ReadBytes((int)block.CompressedSize);
126-
if (_unityCnInfo != null && (block.CompressionType == CompressionType.Lz4 || block.CompressionType == CompressionType.Lz4HC))
127-
{
128-
_unityCnInfo.DecryptBlock(compressedData, compressedData.Length, blockIndex);
129-
}
130-
131-
var uncompressedBuffer = new byte[block.UncompressedSize];
132-
Compression.DecompressToBytes(compressedData, uncompressedBuffer.AsSpan(0, (int)block.UncompressedSize), block.CompressionType);
133-
134-
return uncompressedBuffer;
135-
},
136-
size: block.UncompressedSize
137-
);
198+
var blockBuffer = AssetToBlockCache.ContainsKey(key)
199+
? Cache.GetOrCreate(
200+
key: new(_baseReaderProvider, blockIndex),
201+
factory: () => DecompressBlock(blockIndex),
202+
size: block.UncompressedSize
203+
)
204+
: DecompressBlock(blockIndex);
138205

139206
_currentBlockData = new MemoryStream(blockBuffer);
140207

‎UnityAsset.NET/IO/Stream/FileEntryStream.cs‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,16 @@ namespace UnityAsset.NET.IO.Stream;
44

55
public class FileEntryStreamProvider : IStreamProvider
66
{
7-
private readonly IStreamProvider _streamProvider;
8-
private readonly FileEntry _fileEntry;
7+
public readonly IStreamProvider StreamProvider;
8+
public readonly FileEntry FileEntry;
99

1010
public FileEntryStreamProvider(IStreamProvider streamProvider, FileEntry fileEntry)
1111
{
12-
_streamProvider = streamProvider;
13-
_fileEntry = fileEntry;
12+
StreamProvider = streamProvider;
13+
FileEntry = fileEntry;
1414
}
1515

16-
public System.IO.Stream OpenStream() => new FileEntryStream(_streamProvider, _fileEntry);
16+
public System.IO.Stream OpenStream() => new FileEntryStream(StreamProvider, FileEntry);
1717
}
1818

1919
public class FileEntryStream : System.IO.Stream
@@ -26,13 +26,13 @@ public FileEntryStream(IStreamProvider streamProvider, FileEntry fileEntry)
2626
{
2727
_blockStream = streamProvider.OpenStream();
2828
_fileEntry = fileEntry;
29-
_blockStream.Seek(_fileEntry.Offset, SeekOrigin.Begin);
29+
_blockStream.Seek((long)_fileEntry.Offset, SeekOrigin.Begin);
3030
}
3131

3232
public override bool CanRead => true;
3333
public override bool CanSeek => true;
3434
public override bool CanWrite => false;
35-
public override long Length => _fileEntry.Size;
35+
public override long Length => (long)_fileEntry.Size;
3636

3737
public override long Position
3838
{
@@ -50,7 +50,7 @@ public override int ReadByte()
5050

5151
public override int Read(Span<byte> buffer)
5252
{
53-
long remaining = _fileEntry.Size - _position;
53+
long remaining = (long)_fileEntry.Size - _position;
5454
if (remaining <= 0)
5555
return 0;
5656

@@ -79,7 +79,7 @@ public override long Seek(long offset, SeekOrigin origin)
7979
throw new ArgumentOutOfRangeException(nameof(offset));
8080

8181
_position = newPosition;
82-
_blockStream.Seek(_fileEntry.Offset + newPosition, SeekOrigin.Begin);
82+
_blockStream.Seek((long)_fileEntry.Offset + newPosition, SeekOrigin.Begin);
8383
return _position;
8484
}
8585

0 commit comments

Comments
 (0)