Skip to content

Commit

Permalink
Added Script Objects, Bulk Data export, Save Package Store Manifest
Browse files Browse the repository at this point in the history
  • Loading branch information
Archengius committed Aug 4, 2023
1 parent b698aba commit cbf84f3
Show file tree
Hide file tree
Showing 5 changed files with 200 additions and 31 deletions.
156 changes: 145 additions & 11 deletions Private/CookedAssetWriter.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 )
{
Expand Down Expand Up @@ -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() );
Expand All @@ -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 );
Expand All @@ -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 );
Expand All @@ -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++;
Expand Down Expand Up @@ -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 );
Expand Down Expand Up @@ -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 );
}
}
16 changes: 14 additions & 2 deletions Private/CookedAssetWriter.h
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ struct FAssetSerializationContext
FPackageId PackageId;
FString PackageHeaderFilename;
FPackageMapExportBundleEntry* BundleData;
FIoStoreReader* IoStoreReader;

FPackageFileSummary Summary;
int32 PackageSummaryEndOffset;
Expand Down Expand Up @@ -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 );
Expand All @@ -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 );
};
42 changes: 28 additions & 14 deletions Private/IoStorePackageMap.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 );
}
}
}

Expand All @@ -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 );
}
}

Expand Down Expand Up @@ -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);
Expand All @@ -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("../../../") );
Expand Down Expand Up @@ -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
Expand All @@ -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++;
}
Expand Down Expand Up @@ -346,4 +359,5 @@ void FIoStorePackageMap::ReadExportBundleData( const FPackageId& PackageId, cons
ArcsAr << ExternalArc.ToExportBundleIndex;
}
}
return &PackageData;
}
Loading

0 comments on commit cbf84f3

Please sign in to comment.