Added Script Objects, Bulk Data export, Save Package Store Manifest

Author: Archengius

Private/CookedAssetWriter.cpp

@@ -3,13 +3,16 @@
 #include "CookedAssetWriter.h"
 #include "IoStorePackageMap.h"
 #include "ZenTools.h"
+#include "Dom/JsonObject.h"
 #include "HAL/FileManager.h"
 #include "Misc/Paths.h"
 #include "Serialization/LargeMemoryWriter.h"
 #include "Serialization/MemoryWriter.h"
 #include "UObject/Class.h"
 #include "UObject/Package.h"
 #include "UObject/SoftObjectPath.h"
+#include "Misc/FileHelper.h"
+#include "Serialization/JsonSerializer.h"
 
 FAssetSerializationWriter::FAssetSerializationWriter( FArchive& Ar, FAssetSerializationContext* Context ) : FArchiveProxy( Ar ), Context( Context )
 {
@@ -59,25 +62,101 @@ FCookedAssetWriter::FCookedAssetWriter(const TSharedPtr<FIoStorePackageMap>& InP
 {
 }
 
-void FCookedAssetWriter::WritePackagesFromContainer( const FIoContainerId& ContainerId )
+void FCookedAssetWriter::WritePackagesFromContainer( const TSharedPtr<FIoStoreReader>& Reader )
 {
+	const FIoContainerId ContainerId = Reader->GetContainerId();
 	UE_LOG( LogIoStoreTools, Display, TEXT("Writing asset files for Container %lld"), ContainerId.Value() );
 
 	FPackageContainerMetadata ContainerMetadata;
 	if ( PackageMap->FindPackageContainerMetadata( ContainerId, ContainerMetadata ) )
 	{
 		for ( const FPackageId& PackageId : ContainerMetadata.PackagesInContainer )
 		{
-			WriteSinglePackage( PackageId );
+			WriteSinglePackage( PackageId, false, Reader );
 		}
 		for ( const FPackageId& OptionalPackageId : ContainerMetadata.OptionalPackagesInContainer )
 		{
-			WriteSinglePackage( OptionalPackageId );
+			WriteSinglePackage( OptionalPackageId, true, Reader );
 		}
 	}
 }
 
-void FCookedAssetWriter::WriteSinglePackage( FPackageId PackageId )
+void FCookedAssetWriter::WriteGlobalScriptObjects(const TSharedPtr<FIoStoreReader>& Reader) const
+{
+	TIoStatusOr<FIoBuffer> ScriptObjectsBuffer = Reader->Read(CreateIoChunkId(0, 0, EIoChunkType::ScriptObjects), FIoReadOptions());
+
+	if ( ScriptObjectsBuffer.IsOk() )
+	{
+		const FString ScriptObjectsFilename = FPaths::Combine( RootOutputDir, TEXT("ScriptObjects.bin") );
+		FFileHelper::SaveArrayToFile( TArrayView<const uint8>( ScriptObjectsBuffer.ValueOrDie().Data(), ScriptObjectsBuffer.ValueOrDie().DataSize() ), *ScriptObjectsFilename );
+
+		UE_LOG( LogIoStoreTools, Display, TEXT("Written ScriptObjects chunk to '%s'"), *ScriptObjectsFilename );
+	}
+}
+
+void FCookedAssetWriter::WritePackageStoreManifest() const
+{
+	const FString PackageStoreFilename = RootOutputDir / TEXT("PackageStoreManifest.json");
+	IFileManager::Get().MakeDirectory( *FPaths::GetPath( PackageStoreFilename ), true );
+
+	const TSharedPtr<FJsonObject> RootObject = MakeShared<FJsonObject>();
+
+	TStringBuilder<64> ChunkIdStringBuilder;
+	auto ChunkIdToString = [&ChunkIdStringBuilder](const FIoChunkId& ChunkId)
+	{
+		ChunkIdStringBuilder.Reset();
+		ChunkIdStringBuilder << ChunkId;
+		return *ChunkIdStringBuilder;
+	};
+	
+	TArray<TSharedPtr<FJsonValue>> FilesArray;
+	for ( const TPair<FIoChunkId, FString>& FilePair : ChunkIdToSavedFileMap )
+	{
+		const TSharedPtr<FJsonObject> FileObject = MakeShared<FJsonObject>();
+		FileObject->SetStringField( TEXT("Path"), FilePair.Value );
+		FileObject->SetStringField( TEXT("ChunkId"), ChunkIdToString( FilePair.Key ) );
+
+		FilesArray.Add( MakeShared<FJsonValueObject>( FileObject ) );
+	}
+	RootObject->SetArrayField( TEXT("Files"), FilesArray );
+
+	TArray<TSharedPtr<FJsonValue>> PackagesArray;
+	for ( const TPair<FName, FSavedPackageInfo>& SavedPackageInfo : SavedPackageMap )
+	{
+		const TSharedPtr<FJsonObject> PackageObject = MakeShared<FJsonObject>();
+		PackageObject->SetStringField( TEXT("Name"), SavedPackageInfo.Key.ToString() );
+
+		if ( !SavedPackageInfo.Value.ExportBundleChunks.IsEmpty() )
+		{
+			TArray<TSharedPtr<FJsonValue>> ExportBundleChunkIdsArray;
+			for ( const FIoChunkId& ExportBundleChunkId : SavedPackageInfo.Value.ExportBundleChunks )
+			{
+				ExportBundleChunkIdsArray.Add( MakeShared<FJsonValueString>( ChunkIdToString( ExportBundleChunkId ) ) );
+			}
+			PackageObject->SetArrayField( TEXT("ExportBundleChunkIds"), ExportBundleChunkIdsArray );
+		}
+
+		if ( !SavedPackageInfo.Value.BulkDataChunks.IsEmpty() )
+		{
+			TArray<TSharedPtr<FJsonValue>> BulkDataChunkIdsArray;
+			for ( const FIoChunkId& BulkDataChunkId : SavedPackageInfo.Value.BulkDataChunks )
+			{
+				BulkDataChunkIdsArray.Add( MakeShared<FJsonValueString>( ChunkIdToString( BulkDataChunkId ) ) );
+			}
+			PackageObject->SetArrayField( TEXT("BulkDataChunkIds"), BulkDataChunkIdsArray );
+		}
+		PackagesArray.Add( MakeShared<FJsonValueObject>( PackageObject ) );
+	}
+	RootObject->SetArrayField( TEXT("Packages"), PackagesArray );
+
+	FString ResultJsonString;
+	FJsonSerializer::Serialize( RootObject.ToSharedRef(), TJsonWriterFactory<>::Create( &ResultJsonString ) );
+
+	check( FFileHelper::SaveStringToFile( ResultJsonString, *PackageStoreFilename ) );
+	UE_LOG( LogIoStoreTools, Display, TEXT("Written PackageStore Manifest to '%s'"), *PackageStoreFilename );
+}
+
+void FCookedAssetWriter::WriteSinglePackage( FPackageId PackageId, bool bIsOptionalSegmentPackage, const TSharedPtr<FIoStoreReader>& Reader )
 {
 	FPackageMapExportBundleEntry ExportBundleEntry;
 	checkf( PackageMap->FindExportBundleData( PackageId, ExportBundleEntry ), TEXT("Failed to find export bundle entry for PackageId %lld"), PackageId.ValueForDebugging() );
@@ -93,13 +172,24 @@ void FCookedAssetWriter::WriteSinglePackage( FPackageId PackageId )
 	SerializationContext.PackageId = PackageId;
 	SerializationContext.PackageHeaderFilename = PackageFilename;
 	SerializationContext.BundleData = &ExportBundleEntry;
+	SerializationContext.IoStoreReader = Reader.Get();
+
+	FSavedPackageInfo& SavedPackageInfo = SavedPackageMap.FindOrAdd( SerializationContext.BundleData->PackageName );
+	SavedPackageInfo.ExportBundleChunks.Add( SerializationContext.BundleData->PackageChunkId );
 
 	// Populate package summary, and also process imports and exports
 	ProcessPackageSummaryAndNamesAndExportsAndImports( SerializationContext );
 
 	// Serialize exports into the separate file (event driven loader expects that)
 	{
-		const FString ExportsFilename = FPaths::ChangeExtension( SerializationContext.PackageHeaderFilename, LexToString( EPackageExtension::Exports ) );
+		FString ExtensionString = LexToString( EPackageExtension::Exports );
+		
+		// Optional segment packages have .o prefix before their extensions, e.g.
+		if ( bIsOptionalSegmentPackage )
+		{
+			ExtensionString.InsertAt( 0, TEXT(".o") );
+		}
+		const FString ExportsFilename = FPaths::ChangeExtension( SerializationContext.PackageHeaderFilename, ExtensionString );
 
 		const TUniquePtr<FArchive> ExportsArchive( IFileManager::Get().CreateFileWriter( *ExportsFilename, FILEWRITE_EvenIfReadOnly ) );
 		checkf( ExportsArchive.IsValid(), TEXT("Failed to load exports file '%s'"), *ExportsFilename );
@@ -112,7 +202,17 @@ void FCookedAssetWriter::WriteSinglePackage( FPackageId PackageId )
 	// Serialize package summary and other necessary data into the main asset header file
 	{
 		const EPackageExtension HeaderExtension = ( SerializationContext.Summary.GetPackageFlags() & PKG_ContainsMap ) != 0 ? EPackageExtension::Map : EPackageExtension::Asset;
-		const FString HeaderFilename = FPaths::ChangeExtension( SerializationContext.PackageHeaderFilename, LexToString( HeaderExtension ) );
+		FString ExtensionString = LexToString( HeaderExtension );
+
+		// Optional segment packages have .o prefix before their extensions, e.g.
+		if ( bIsOptionalSegmentPackage )
+		{
+			ExtensionString.InsertAt( 0, TEXT(".o") );
+		}
+		const FString HeaderFilename = FPaths::ChangeExtension( SerializationContext.PackageHeaderFilename, ExtensionString );
+		
+		FString RelativeFilename = FPaths::SetExtension( ExportBundleEntry.PackageFilename, ExtensionString );
+		ChunkIdToSavedFileMap.Add( SerializationContext.BundleData->PackageChunkId, RelativeFilename );
 
 		const TUniquePtr<FArchive> HeaderArchive( IFileManager::Get().CreateFileWriter( *HeaderFilename, FILEWRITE_EvenIfReadOnly ) );
 		checkf( HeaderArchive.IsValid(), TEXT("Failed to open header file '%s'"), *HeaderFilename );
@@ -122,6 +222,9 @@ void FCookedAssetWriter::WriteSinglePackage( FPackageId PackageId )
 		HeaderArchive->Flush();
 	}
 
+	// Write bulk data
+	WriteBulkData( SerializationContext );
+
 	// Notify the user that we have finished writing the asset
 	UE_LOG( LogIoStoreTools, Display, TEXT("Serialized Package '%s' to '%s'"), *SerializationContext.BundleData->PackageName.ToString(), *SerializationContext.PackageHeaderFilename );
 	NumPackagesWritten++;
@@ -584,8 +687,8 @@ FPackageIndex FCookedAssetWriter::CreateObjectExport( const FPackageMapExportEnt
 
 	ObjectExport.ObjectFlags = ExportData.ObjectFlags;
 
-	ObjectExport.SerialSize = ExportData.CookedSerialData->Num();
-	ObjectExport.SerialOffset = -1; // Not resolved yet
+	ObjectExport.SerialSize = INDEX_NONE;
+	ObjectExport.SerialOffset = INDEX_NONE;
 
 	ObjectExport.bForcedExport = false; // not serialized
 	ObjectExport.bNotForClient = EnumHasAnyFlags( ExportData.FilterFlags, EExportFilterFlags::NotForClient );
@@ -868,20 +971,51 @@ void FCookedAssetWriter::WritePackageHeader(FArchive& Ar, FAssetSerializationCon
 
 void FCookedAssetWriter::WritePackageExports(FArchive& Ar, FAssetSerializationContext& Context)
 {
+	// Open the package bundle chunk to read exports
+	TIoStatusOr<FIoBuffer> ChunkBuffer = Context.IoStoreReader->Read( Context.BundleData->PackageChunkId, FIoReadOptions() );
+	check( ChunkBuffer.IsOk() );
+	const uint8* ChunkDataStart = ChunkBuffer.ValueOrDie().Data();
+	const uint8* ChunkDataEnd = ChunkDataStart + ChunkBuffer.ValueOrDie().DataSize();
+	
 	// Write export blobs
 	for ( int32 i = 0; i < Context.ExportMap.Num(); i++ )
 	{
-		TArray<uint8>& SerialData = *Context.BundleData->ExportMap[ i ].CookedSerialData;
+		const FPackageMapExportEntry& OriginalExport = Context.BundleData->ExportMap[ i ];
 		FObjectExport& Export = Context.ExportMap[ i ];
 
 		Export.SerialOffset = Ar.Tell();
-		Export.SerialSize = SerialData.Num();
+		Export.SerialSize = OriginalExport.SerialDataSize;
 
-		Ar.Serialize( SerialData.GetData(), SerialData.Num() );
+		const uint8* SerialDataStart = ChunkDataStart + OriginalExport.SerialDataOffset;
+		check( SerialDataStart <= ChunkDataEnd );
+		Ar.Serialize( const_cast<uint8*>( SerialDataStart ), OriginalExport.SerialDataSize );
 	}
 	Context.Summary.BulkDataStartOffset = Ar.Tell();
 
 	// Exports end with the package file tag
 	uint32 FooterData = PACKAGE_FILE_TAG;
 	Ar << FooterData;
 }
+
+void FCookedAssetWriter::WriteBulkData( const FAssetSerializationContext& Context )
+{
+	FSavedPackageInfo& SavedPackageInfo = SavedPackageMap.FindOrAdd( Context.BundleData->PackageName );
+	
+	for ( const FIoChunkId& BulkDataChunkId : Context.BundleData->BulkDataChunkIds )
+	{
+		TIoStatusOr<FIoBuffer> BulkDataBuffer = Context.IoStoreReader->Read( BulkDataChunkId, FIoReadOptions() );
+		check( BulkDataBuffer.IsOk() );
+
+		TIoStatusOr<FIoStoreTocChunkInfo> ChunkInfo = Context.IoStoreReader->GetChunkInfo( BulkDataChunkId );
+		check( ChunkInfo.IsOk() );
+
+		FString RelativeFilename = ChunkInfo.ValueOrDie().FileName;
+		RelativeFilename.RemoveFromStart( TEXT("../../../") );
+
+		const FString ResultFilename = FPaths::Combine( RootOutputDir, RelativeFilename );
+		FFileHelper::SaveArrayToFile( TArrayView<const uint8>( BulkDataBuffer.ValueOrDie().Data(), BulkDataBuffer.ValueOrDie().DataSize() ), *ResultFilename );
+
+		ChunkIdToSavedFileMap.Add( BulkDataChunkId, RelativeFilename );
+		SavedPackageInfo.BulkDataChunks.Add( BulkDataChunkId );
+	}
+}

Private/CookedAssetWriter.h

@@ -36,6 +36,7 @@ struct FAssetSerializationContext
 	FPackageId PackageId;
 	FString PackageHeaderFilename;
 	FPackageMapExportBundleEntry* BundleData;
+	FIoStoreReader* IoStoreReader;
 
 	FPackageFileSummary Summary;
 	int32 PackageSummaryEndOffset;
@@ -64,20 +65,30 @@ class FAssetSerializationWriter : public FArchiveProxy
 	virtual void SetFilterEditorOnly(bool InFilterEditorOnly) override;
 };
 
+struct FSavedPackageInfo
+{
+	TArray<FIoChunkId> ExportBundleChunks;
+	TArray<FIoChunkId> BulkDataChunks;
+};
+
 class ZENTOOLS_API FCookedAssetWriter
 {
 protected:
 	TSharedPtr<FIoStorePackageMap> PackageMap;
 	FString RootOutputDir;
 	int32 NumPackagesWritten;
+	TMap<FIoChunkId, FString> ChunkIdToSavedFileMap;
+	TMap<FName, FSavedPackageInfo> SavedPackageMap;
 public:
 	FCookedAssetWriter( const TSharedPtr<FIoStorePackageMap>& InPackageMap, const FString& InOutputDir );
 
-	void WritePackagesFromContainer( const FIoContainerId& ContainerId );
+	void WritePackagesFromContainer( const TSharedPtr<FIoStoreReader>& Reader );
+	void WriteGlobalScriptObjects( const TSharedPtr<FIoStoreReader>& Reader ) const;
+	void WritePackageStoreManifest() const;
 
 	FORCEINLINE int32 GetTotalNumPackagesWritten() const { return NumPackagesWritten; }
 private:
-	void WriteSinglePackage( FPackageId PackageId );
+	void WriteSinglePackage( FPackageId PackageId, bool bIsOptionalSegmentPackage, const TSharedPtr<FIoStoreReader>& Reader );
 	void ProcessPackageSummaryAndNamesAndExportsAndImports( FAssetSerializationContext& Context ) const;
 	static FExportBundleEntry BuildPreloadDependenciesFromExportBundle( int32 ExportBundleIndex, FAssetSerializationContext& Context );
 	static void BuildPreloadDependenciesFromArcs( FAssetSerializationContext& Context );
@@ -98,4 +109,5 @@ class ZENTOOLS_API FCookedAssetWriter
 
 	static void WritePackageHeader( FArchive& Ar, FAssetSerializationContext& Context );
 	static void WritePackageExports( FArchive& Ar, FAssetSerializationContext& Context );
+	void WriteBulkData(const FAssetSerializationContext& Context );
 };

Private/IoStorePackageMap.cpp

@@ -65,10 +65,20 @@ void FIoStorePackageMap::PopulateFromContainer(const TSharedPtr<FIoStoreReader>&
 
 		TIoStatusOr<FIoStoreTocChunkInfo> ChunkInfo = Reader->GetChunkInfo( ChunkId );
 		TIoStatusOr<FIoBuffer> PackageBuffer = Reader->Read( ChunkId, FIoReadOptions() );
-		
-		if ( PackageBuffer.IsOk() )
+		check( PackageBuffer.IsOk() );
+
+		FPackageMapExportBundleEntry* ExportBundleEntry = ReadExportBundleData( PackageId, ChunkInfo.ValueOrDie(), PackageBuffer.ValueOrDie() );
+
+		// Required segment packages can have bulk data, memory mapped bulk data and optional bulk data
+		const TArray BulkDataChunkTypes{ EIoChunkType::BulkData, EIoChunkType::MemoryMappedBulkData, EIoChunkType::OptionalBulkData };
+
+		for ( const EIoChunkType BulkDataChunkType : BulkDataChunkTypes )
 		{
-			ReadExportBundleData( PackageId, ChunkInfo.ValueOrDie(), PackageBuffer.ValueOrDie() );
+			const FIoChunkId BulkDataChunkId = CreateIoChunkId( PackageId.Value(), 0, BulkDataChunkType );
+			if ( Reader->GetChunkInfo( BulkDataChunkId ).IsOk() )
+			{
+				ExportBundleEntry->BulkDataChunkIds.Add( BulkDataChunkId );
+			}
 		}
 	}
 
@@ -80,10 +90,15 @@ void FIoStorePackageMap::PopulateFromContainer(const TSharedPtr<FIoStoreReader>&
 
 		TIoStatusOr<FIoStoreTocChunkInfo> ChunkInfo = Reader->GetChunkInfo( ChunkId );
 		TIoStatusOr<FIoBuffer> PackageBuffer = Reader->Read( ChunkId, FIoReadOptions() );
+		check( PackageBuffer.IsOk() );
 
-		if ( PackageBuffer.IsOk() )
+		FPackageMapExportBundleEntry* ExportBundleEntry = ReadExportBundleData( PackageId, ChunkInfo.ValueOrDie(), PackageBuffer.ValueOrDie() );
+
+		// Optional segment packages can only have optional segment bulk data
+		const FIoChunkId BulkDataChunkId = CreateIoChunkId( PackageId.Value(), 1, EIoChunkType::BulkData );
+		if ( Reader->GetChunkInfo( BulkDataChunkId ).IsOk() )
 		{
-			ReadExportBundleData( PackageId, ChunkInfo.ValueOrDie(), PackageBuffer.ValueOrDie() );
+			ExportBundleEntry->BulkDataChunkIds.Add( BulkDataChunkId );
 		}
 	}
 
@@ -195,7 +210,7 @@ FPackageLocalObjectRef FIoStorePackageMap::ResolvePackageLocalRef( const FPackag
 	return Result;
 }
 
-void FIoStorePackageMap::ReadExportBundleData( const FPackageId& PackageId, const FIoStoreTocChunkInfo& ChunkInfo, const FIoBuffer& ChunkBuffer )
+FPackageMapExportBundleEntry* FIoStorePackageMap::ReadExportBundleData( const FPackageId& PackageId, const FIoStoreTocChunkInfo& ChunkInfo, const FIoBuffer& ChunkBuffer )
 {
 	const uint8* PackageSummaryData = ChunkBuffer.Data();
 	const FZenPackageSummary* PackageSummary = reinterpret_cast<const FZenPackageSummary*>(PackageSummaryData);
@@ -221,6 +236,7 @@ void FIoStorePackageMap::ReadExportBundleData( const FPackageId& PackageId, cons
 	PackageData.PackageName = PackageName;
 	PackageData.PackageFlags = PackageSummary->PackageFlags;
 	PackageData.VersioningInfo = VersioningInfo;
+	PackageData.PackageChunkId = ChunkInfo.Id;
 
 	// get rid of standard filename prefix
 	PackageData.PackageFilename.RemoveFromStart( TEXT("../../../") );
@@ -279,8 +295,8 @@ void FIoStorePackageMap::ReadExportBundleData( const FPackageId& PackageId, cons
 		ExportData.SuperIndex = ResolvePackageLocalRef( ExportMapEntry.SuperIndex, PackageHeader.ImportedPackages, ImportedPublicExportHashes );
 		ExportData.TemplateIndex = ResolvePackageLocalRef( ExportMapEntry.TemplateIndex, PackageHeader.ImportedPackages, ImportedPublicExportHashes );
 
-		ExportData.CookedSerialData = MakeShared<TArray<uint8>>();
-		ExportData.CookedSerialData->SetNumUninitialized(ExportMapEntry.CookedSerialSize);
+		ExportData.SerialDataSize = ExportMapEntry.CookedSerialSize;
+		ExportData.SerialDataOffset = INDEX_NONE;
 	}
 
 	// Read export bundles
@@ -303,12 +319,9 @@ void FIoStorePackageMap::ReadExportBundleData( const FPackageId& PackageId, cons
 
 			if (BundleEntry->CommandType == FExportBundleEntry::ExportCommandType_Serialize)
 			{
-				const FPackageMapExportEntry& Export = PackageData.ExportMap[ BundleEntry->LocalExportIndex ];
-				const int32 ExportSerialSize = Export.CookedSerialData->Num();
-
-				// Copy the export data into the buffer that we have crated earlier
-				FMemory::Memcpy( Export.CookedSerialData->GetData(), PackageSummaryData + CurrentExportOffset, ExportSerialSize );
-				CurrentExportOffset += ExportSerialSize;
+				FPackageMapExportEntry& Export = PackageData.ExportMap[ BundleEntry->LocalExportIndex ];
+				Export.SerialDataOffset = CurrentExportOffset;
+				CurrentExportOffset += Export.SerialDataSize;
 			}
 			BundleEntry++;
 		}
@@ -346,4 +359,5 @@ void FIoStorePackageMap::ReadExportBundleData( const FPackageId& PackageId, cons
 			ArcsAr << ExternalArc.ToExportBundleIndex;
 		}
 	}
+	return &PackageData;
 }