11using System . Collections . Concurrent ;
2+ using System . Collections . Frozen ;
23using System . Diagnostics ;
34using AssetRipper . Tpk ;
45using AssetRipper . Tpk . TypeTrees ;
89using UnityAsset . NET . Files . SerializedFiles ;
910using UnityAsset . NET . FileSystem ;
1011using UnityAsset . NET . IO ;
11- using UnityAsset . NET . IO . Stream ;
12+ using UnityAsset . NET . IO . Reader ;
1213using UnityAsset . NET . TypeTree ;
1314using UnityAsset . NET . TypeTree . PreDefined . Types ;
1415using 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}
0 commit comments