Skip to content

Commit e232b6c

Browse files
committed
IO Refactor
1 parent dd8da93 commit e232b6c

35 files changed

Lines changed: 1091 additions & 865 deletions

UnityAsset.NET.TypeTreeHelper/TypeTreeRepr.cs

Lines changed: 2 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ public class TypeTreeRepr : IEquatable<TypeTreeRepr>
1111
public readonly TypeTreeRepr[] SubNodes;
1212

1313
public bool RequiresAlign { get; }
14-
14+
public bool IsNamed { get; }
1515
public int Hash { get; }
1616

1717
private TypeTreeRepr(string name, string typeName, TypeTreeRepr[] subNodes, bool requiresAlign)
@@ -20,6 +20,7 @@ private TypeTreeRepr(string name, string typeName, TypeTreeRepr[] subNodes, bool
2020
TypeName = typeName;
2121
SubNodes = subNodes;
2222
RequiresAlign = requiresAlign;
23+
IsNamed = subNodes.Any(x => x.Name == "m_Name");
2324

2425
if (TypeName == "map")
2526
RequiresAlign |= SubNodes[0].SubNodes[1].RequiresAlign;
@@ -85,16 +86,4 @@ public bool Equals(TypeTreeRepr? other)
8586

8687
return true;
8788
}
88-
89-
public static bool operator ==(TypeTreeRepr? left, TypeTreeRepr? right)
90-
{
91-
if (left is null)
92-
return right is null;
93-
return left.Equals(right);
94-
}
95-
96-
public static bool operator !=(TypeTreeRepr? left, TypeTreeRepr? right)
97-
{
98-
return !(left == right);
99-
}
10089
}

UnityAsset.NET/Asset.cs

Lines changed: 45 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,14 @@
11
using UnityAsset.NET.Files.SerializedFiles;
22
using UnityAsset.NET.IO.Reader;
3-
using UnityAsset.NET.IO.Stream;
43
using UnityAsset.NET.TypeTree;
54
using UnityAsset.NET.TypeTree.PreDefined;
65

76
namespace UnityAsset.NET;
87

9-
public class Asset
8+
public class Asset : IEquatable<Asset>
109
{
1110
public readonly AssetFileInfo Info;
12-
private IUnityAsset? _value;
11+
private WeakReference<IUnityAsset>? _value;
1312
private string? _name;
1413
private readonly Lock _lock = new();
1514
public bool IsNamedAsset;
@@ -29,23 +28,39 @@ private AssetReader DataReader
2928

3029
public SerializedFile SourceFile { get; }
3130

31+
private IUnityAsset GetValue()
32+
{
33+
var value = UnityObjectFactory.Create(Info.Type, DataReader);
34+
BlockReader.OnAssetParsed(this);
35+
return value;
36+
}
37+
38+
3239
public IUnityAsset Value
3340
{
3441
get
3542
{
3643
lock (_lock)
3744
{
38-
if (_value == null)
45+
if (_value is null)
3946
{
40-
using var reader = DataReader;
41-
_value = UnityObjectFactory.Create(Info.Type, reader);
42-
BlockStream.OnAssetParsed(this);
47+
var value = GetValue();
48+
_value = new WeakReference<IUnityAsset>(value);
4349
if (IsNamedAsset)
4450
{
45-
_name = ((INamedObject)_value).m_Name;
51+
_name = ((INamedObject)value).m_Name;
4652
}
53+
54+
return value;
4755
}
48-
return _value;
56+
else if (_value.TryGetTarget(out var value))
57+
{
58+
return value;
59+
}
60+
61+
var newValue = GetValue();
62+
_value = new WeakReference<IUnityAsset>(newValue);
63+
return newValue;
4964
}
5065
}
5166
}
@@ -88,16 +103,31 @@ public string Container
88103
public Asset(SerializedFile sf, AssetFileInfo info)
89104
{
90105
Info = info;
91-
IsNamedAsset = info.Type.Nodes.Any(n => n is {Name: "m_Name", Type: "string", Level: 1} );
106+
IsNamedAsset = info.Type.IsNamed;
92107

93108
SourceFile = sf;
94109
}
110+
111+
112+
public bool Equals(Asset? other)
113+
{
114+
if (ReferenceEquals(this, other))
115+
return true;
116+
117+
if (other is null)
118+
return false;
119+
120+
return ReferenceEquals(SourceFile, other.SourceFile)
121+
&& PathId == other.PathId;
122+
}
123+
124+
public override bool Equals(object? obj)
125+
=> obj is Asset other && Equals(other);
95126

96-
public void Release()
127+
public override int GetHashCode()
97128
{
98-
lock (_lock)
99-
{
100-
_value = null;
101-
}
129+
return HashCode.Combine(
130+
SourceFile,
131+
PathId);
102132
}
103133
}

UnityAsset.NET/AssetHelper/SpriteHelper.cs

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -162,9 +162,8 @@ private static Vector2f[][] GetTriangles(ISpriteRenderData m_RD, UnityRevision v
162162
// TODO: m_Streams Interface for IVertexData
163163
var m_Streams = m_VertexData.GetStreams(version);
164164
var m_Stream = m_Streams[m_Channel.stream];
165-
using var vertexReader =
166-
new CustomStreamReader(new MemoryStream(m_VertexData.m_DataSize.data), Endianness.LittleEndian);
167-
using var indexReader = new CustomStreamReader(new MemoryStream(m_RD.m_IndexBuffer.ToArray()), Endianness.LittleEndian);
165+
var vertexReader = new MemoryReader(m_VertexData.m_DataSize.data, endian: Endianness.LittleEndian);
166+
var indexReader = new MemoryReader(m_RD.m_IndexBuffer, endian: Endianness.LittleEndian);
168167
foreach (var subMesh in m_RD.m_SubMeshes)
169168
{
170169
vertexReader.Position = m_Stream.offset + subMesh.firstVertex * m_Stream.stride + m_Channel.offset;

UnityAsset.NET/AssetManager.cs

Lines changed: 31 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
using System.Collections.Concurrent;
2+
using System.Collections.Frozen;
23
using System.Diagnostics;
34
using AssetRipper.Tpk;
45
using AssetRipper.Tpk.TypeTrees;
@@ -8,7 +9,7 @@
89
using UnityAsset.NET.Files.SerializedFiles;
910
using UnityAsset.NET.FileSystem;
1011
using UnityAsset.NET.IO;
11-
using UnityAsset.NET.IO.Stream;
12+
using UnityAsset.NET.IO.Reader;
1213
using UnityAsset.NET.TypeTree;
1314
using UnityAsset.NET.TypeTree.PreDefined.Types;
1415
using UnityAsset.NET.TypeTreeHelper;
@@ -19,16 +20,17 @@ public class AssetManager
1920
{
2021
private IFileSystem _fileSystem;
2122

22-
public ConcurrentDictionary<string, IFile> LoadedFiles = new();
23+
public FrozenDictionary<string, IFile> LoadedFiles = new Dictionary<string, IFile>().ToFrozenDictionary();
2324

24-
public ConcurrentDictionary<IVirtualFile, IFile> VirtualFileToFileMap = new();
25+
public ConcurrentDictionary<IVirtualFileInfo, IFile> VirtualFileToFileMap = new();
2526

2627
public UnityRevision? Version { get; private set; }
2728
public BuildTarget? BuildTarget { get; private set; }
2829

29-
private List<SerializedType> _loadedTypes = new();
3030
public List<Asset> LoadedAssets { get; private set; } = new();
3131

32+
public static Dictionary<Hash128, TypeTreeRepr> LoadedTypes = new();
33+
3234
public AssetManager(IFileSystem? fileSystem = null, IFileSystem.ErrorHandler? onError = null)
3335
{
3436
_fileSystem = fileSystem ?? new FileSystem.DirectFileSystem.DirectFileSystem(onError);
@@ -40,12 +42,11 @@ public void SetFileSystem(IFileSystem fileSystem)
4042
_fileSystem = fileSystem;
4143
}
4244

43-
public async Task LoadAsync(List<IVirtualFile> files, bool ignoreDuplicatedFiles = false, IProgress<LoadProgress>? progress = null)
45+
public async Task LoadAsync(List<IVirtualFileInfo> files, bool ignoreDuplicatedFiles = false, IProgress<LoadProgress>? progress = null)
4446
{
4547
await Task.Run(() =>
4648
{
4749
var fileWrappers = new ConcurrentBag<(string, IFile)>();
48-
var types = new ConcurrentBag<SerializedType>();
4950
int progressCount = 0;
5051
var total = files.Count;
5152

@@ -57,18 +58,15 @@ await Task.Run(() =>
5758
{
5859
case FileType.BundleFile:
5960
{
60-
var bundleFile = new BundleFile(file);
61-
bundleFile.ParseFilesWithTypeConversion();
61+
var bundleFile = new BundleFile(file, lazyLoad: false);
6262
foreach (var fw in bundleFile.Files)
6363
{
6464
fileWrappers.Add((fw.Info.Path, fw.File));
6565

6666
if (fw is { File: SerializedFile sf })
6767
{
68-
if (!sf.Metadata.TypeTreeEnabled)
68+
if (!sf.Metadata.TypeTreeEnabled && !anyTypeTreeDisabled)
6969
Volatile.Write(ref anyTypeTreeDisabled, true);
70-
foreach(var type in sf.Metadata.Types)
71-
types.Add(type);
7270
}
7371
}
7472

@@ -82,9 +80,6 @@ await Task.Run(() =>
8280

8381
if (!serializedFile.Metadata.TypeTreeEnabled)
8482
Volatile.Write(ref anyTypeTreeDisabled, true);
85-
86-
foreach(var type in serializedFile.Metadata.Types)
87-
types.Add(type);
8883

8984
VirtualFileToFileMap[file] = serializedFile;
9085
break;
@@ -94,14 +89,21 @@ await Task.Run(() =>
9489
progress?.Report(new LoadProgress($"AssetManager: Loading {file.Name}", total, currentProgress));
9590
});
9691

92+
BlockReader.RemoveSingleReferenceBlocks();
93+
BlockReader.Cache = new(maxSize: BlockReader.TotalBlockSize * 3 / 4); // it works good for BuildSceneHierarchy
94+
95+
var tmpLoadedFilesDict = new Dictionary<string, IFile>();
96+
9797
foreach (var (path, file) in fileWrappers)
9898
{
99-
if (!LoadedFiles.TryAdd(path, file) && !ignoreDuplicatedFiles)
99+
if (!tmpLoadedFilesDict.TryAdd(path, file) && !ignoreDuplicatedFiles)
100100
{
101101
throw new InvalidOperationException($"File {path} already loaded");
102102
}
103103
}
104104

105+
LoadedFiles = tmpLoadedFilesDict.ToFrozenDictionary();
106+
105107
var firstFile = LoadedFiles.Values
106108
.FirstOrDefault(file => file is SerializedFile);
107109
if (firstFile is SerializedFile firstSerializedFile)
@@ -110,15 +112,6 @@ await Task.Run(() =>
110112
Version = firstSerializedFile.Metadata.UnityVersion;
111113
}
112114

113-
var seenHashes = new HashSet<Hash128>();
114-
foreach (var type in types)
115-
{
116-
if (seenHashes.Add(type.TypeHash))
117-
{
118-
_loadedTypes.Add(type);
119-
}
120-
}
121-
122115
BuildUnityTypes(anyTypeTreeDisabled, progress);
123116

124117
LoadedAssets = LoadedFiles.Values
@@ -151,24 +144,21 @@ private void BuildUnityTypes(bool anyTypeTreeDisabled, IProgress<LoadProgress>?
151144

152145
var rootTypeNodesMap = anyTypeTreeDisabled ? TpkUnityTreeNodeFactory.GetRootTypeNodes(Version!.ToString()) : null;
153146

154-
var rootTypes = new Dictionary<Hash128, TypeTreeRepr>();
155-
156-
foreach (var type in _loadedTypes)
147+
foreach (var (hash, (nodes, typeId, repr)) in TypeTreeNode.Cache)
157148
{
158-
var nodes = type.Nodes;
159-
if (nodes.Count == 0 && anyTypeTreeDisabled)
149+
if (nodes.Length == 0 && anyTypeTreeDisabled)
160150
{
161-
var typeName = type.ToTypeName();
162-
rootTypes.Add(type.TypeHash, rootTypeNodesMap![typeName]);
151+
var typeName = ((AssetClassID)typeId).ToString();
152+
LoadedTypes.Add(hash, rootTypeNodesMap![typeName]);
163153
}
164154
else
165155
{
166-
rootTypes.Add(type.TypeHash, nodes[0].ToTypeTreeRepr(nodes));
156+
LoadedTypes.Add(hash, repr ?? nodes[0].ToTypeTreeRepr(nodes));
167157
}
168158
}
169159

170-
AssemblyManager.LoadTypes(rootTypes);
171-
progress?.Report(new LoadProgress($"AssetManager: Generated {_loadedTypes.Count} types", 2, 2));
160+
AssemblyManager.LoadTypes(LoadedTypes);
161+
progress?.Report(new LoadProgress($"AssetManager: Generated {TypeTreeNode.Cache.Count} types", 2, 2));
172162

173163
foreach (var (_, file) in LoadedFiles)
174164
{
@@ -185,7 +175,7 @@ private void BuildUnityTypes(bool anyTypeTreeDisabled, IProgress<LoadProgress>?
185175
foreach (var asset in sf.Assets)
186176
{
187177
if (asset.IsNamedAsset) continue;
188-
var typeTreeRepr = rootTypes[asset.Info.Type.TypeHash];
178+
var typeTreeRepr = LoadedTypes[asset.Info.Type.TypeHash];
189179
foreach (var field in typeTreeRepr.SubNodes)
190180
{
191181
if (field.Name == "m_Name")
@@ -218,8 +208,9 @@ public async Task LoadDirectoryAsync(string directoryPath, bool ignoreDuplicated
218208
var path = streamingInfo.path.Split('/')[^1];
219209
if (LoadedFiles.TryGetValue(path, out IFile? file))
220210
{
221-
if (file is IReader reader)
211+
if (file is IReaderProvider readerProvider)
222212
{
213+
var reader = readerProvider.CreateReader();
223214
reader.Seek((long)streamingInfo.offset);
224215
return reader.ReadBytes((int)streamingInfo.size);
225216
}
@@ -231,9 +222,9 @@ public async Task LoadDirectoryAsync(string directoryPath, bool ignoreDuplicated
231222
public void Clear()
232223
{
233224
LoadedAssets = new();
234-
_loadedTypes = new();
225+
LoadedTypes = new();
235226
VirtualFileToFileMap = new();
236-
LoadedFiles = new();
227+
LoadedFiles = new Dictionary<string, IFile>().ToFrozenDictionary();
237228
Version = null;
238229
BuildTarget = null;
239230

@@ -243,7 +234,7 @@ public void Clear()
243234
TypeTreeRepr.Cache = new();
244235
TpkUnityTreeNodeFactory.Deinit();
245236
AssemblyManager.CleanCache();
246-
BlockStream.Cache.Reset(Setting.DefaultBlockCacheSize);
247-
BlockStream.AssetToBlockCache = new();
237+
BlockReader.Cache.Reset(Setting.DefaultBlockCacheSize);
238+
BlockReader.AssetToBlockCache = new();
248239
}
249240
}

UnityAsset.NET/Extensions/SerializedFileExtensions.cs

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,15 +13,13 @@ public static void ProcessAssetBundle(this SerializedFile sf)
1313
{
1414
var assetBundle = (IAssetBundle)asset.Value;
1515

16-
sf.PreloadTable = assetBundle.m_PreloadTable;
17-
1816
foreach (var (container, assetInfo) in assetBundle.m_Container)
1917
{
2018
var preloadIndex = assetInfo.preloadIndex;
2119
var preloadSize = assetInfo.preloadSize;
2220
for (int i = preloadIndex; i < preloadIndex + preloadSize; i++)
2321
{
24-
var pptr = sf.PreloadTable[i];
22+
var pptr = assetBundle.m_PreloadTable[i];
2523
sf.Containers[pptr.m_PathID] = container;
2624
}
2725
}

UnityAsset.NET/Extensions/TypeTreeNodeExtensions.cs

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,23 +20,22 @@ public static TypeTreeNode Parent(this TypeTreeNode current, List<TypeTreeNode>
2020
throw new IndexOutOfRangeException();
2121
}
2222

23-
public static List<TypeTreeNode> Children(this TypeTreeNode current, List<TypeTreeNode> nodes)
23+
public static List<TypeTreeNode> Children(this TypeTreeNode current, TypeTreeNode[] nodes)
2424
{
25-
var nodesSpan = nodes.AsReadOnlySpan();
26-
if ((int)current.Index + 1 >= nodesSpan.Length || nodesSpan[(int)current.Index + 1].Level <= current.Level)
25+
if ((int)current.Index + 1 >= nodes.Length || nodes[(int)current.Index + 1].Level <= current.Level)
2726
return new();
2827
var children = new List<TypeTreeNode>();
29-
for (int i = (int)current.Index + 1; i < nodesSpan.Length; i++)
28+
for (int i = (int)current.Index + 1; i < nodes.Length; i++)
3029
{
31-
if (nodesSpan[i].Level <= current.Level)
30+
if (nodes[i].Level <= current.Level)
3231
break;
33-
if (nodesSpan[i].Level == current.Level + 1)
34-
children.Add(nodesSpan[i]);
32+
if (nodes[i].Level == current.Level + 1)
33+
children.Add(nodes[i]);
3534
}
3635
return children;
3736
}
3837

39-
public static TypeTreeRepr ToTypeTreeRepr(this TypeTreeNode current, List<TypeTreeNode> nodes)
38+
public static TypeTreeRepr ToTypeTreeRepr(this TypeTreeNode current, TypeTreeNode[] nodes)
4039
{
4140
if (String.IsNullOrEmpty(current.Type) || String.IsNullOrEmpty(current.Name))
4241
throw new Exception("Type/Name is empty");

0 commit comments

Comments
 (0)