UShaderLab/Source/UShaderLab/Private/ShaderLabMaterialRegistry.cpp

// Copyright UShaderLab. All Rights Reserved.

#include "ShaderLabMaterialRegistry.h"

#include "Materials/Material.h"
#include "Misc/Paths.h"
#include "ShaderCore.h"
#include "ShaderLabMaterialAssetUserData.h"
#include "ShaderLabModel.h"
#include "Serialization/AsyncLoadingEvents.h"
#include "UObject/Package.h"
#include "UObject/UObjectGlobals.h"

DEFINE_LOG_CATEGORY_STATIC(LogShaderLabRegistry, Log, All);

namespace
{
	const TCHAR* GShaderLabPackageName = TEXT("/Script/UShaderLab");
}

FShaderLabMaterialRegistry& FShaderLabMaterialRegistry::Get()
{
	static FShaderLabMaterialRegistry Instance;
	return Instance;
}

UPackage* FShaderLabMaterialRegistry::GetPackage()
{
	if (!Package)
	{
		// Mirror the engine's native-package convention (and AngelScript literal assets):
		// a compiled-in /Script package whose objects resolve via FindObject, never disk.
		Package = NewObject<UPackage>(nullptr, GShaderLabPackageName, RF_Public | RF_Standalone | RF_MarkAsRootSet);
		Package->SetPackageFlags(PKG_CompiledIn);
		Package->AddToRoot();
	}
	return Package;
}

FName FShaderLabMaterialRegistry::MakeObjectName(const FString& ShaderName)
{
	FString Out;
	Out.Reserve(ShaderName.Len());
	for (const TCHAR C : ShaderName)
	{
		Out.AppendChar((FChar::IsAlnum(C) || C == TEXT('_')) ? C : TEXT('_'));
	}
	if (Out.IsEmpty())
	{
		Out = TEXT("Unnamed");
	}
	return FName(*Out);
}

FString FShaderLabMaterialRegistry::MakeObjectPath(const FString& ShaderName)
{
	return FString::Printf(TEXT("%s.%s"), GShaderLabPackageName, *MakeObjectName(ShaderName).ToString());
}

FString FShaderLabMaterialRegistry::MakeVirtualShaderPath(const FString& AbsoluteSourcePath)
{
	FString Full = FPaths::ConvertRelativePathToFull(AbsoluteSourcePath);
	Full.ReplaceInline(TEXT("\\"), TEXT("/"));

	// Reverse the registered shader-source mappings (/Project, /Plugin/<AnyPlugin>, /Engine, ...) —
	// generic, so any plugin that authors `.usl` under a mapped Shaders root gets a stable virtual path,
	// with no hardcoded plugin name. Pick the longest matching real dir (most specific root); on a tie,
	// the lexicographically smaller virtual key for determinism. The mappings are editor/cook-populated;
	// at cooked runtime the map is empty and the ShaderLabPath is unused, so the absolute-path fallback is fine.
	FString BestVirtual;
	FString BestKey;
	int32 BestRealLen = -1;
	for (const TPair<FString, FString>& Mapping : AllShaderSourceDirectoryMappings())
	{
		FString Real = FPaths::ConvertRelativePathToFull(Mapping.Value);
		Real.ReplaceInline(TEXT("\\"), TEXT("/"));
		if (!Full.StartsWith(Real + TEXT("/"), ESearchCase::IgnoreCase))
		{
			continue;
		}
		const bool bBetter = (Real.Len() > BestRealLen)
			|| (Real.Len() == BestRealLen && Mapping.Key < BestKey);
		if (bBetter)
		{
			BestRealLen = Real.Len();
			BestKey = Mapping.Key;
			// Full.RightChop(Real.Len()) begins with '/', so the result is "<Key>/<rel>".
			BestVirtual = Mapping.Key + Full.RightChop(Real.Len());
		}
	}
	return (BestRealLen >= 0) ? BestVirtual : Full;
}

UMaterial* FShaderLabMaterialRegistry::FindMaterial(const FString& ShaderName) const
{
	if (!Package)
	{
		return nullptr;
	}
	return FindObject<UMaterial>(Package, *MakeObjectName(ShaderName).ToString());
}

UMaterial* FShaderLabMaterialRegistry::RegisterFromModel(const FShaderLabModel& Model)
{
	UPackage* Pkg = GetPackage();
	const FName ObjName = MakeObjectName(Model.ShaderName);

	// Contract: a material's identity is its file name. Two different source files that sanitize to the
	// same object name are a hard error (would otherwise silently clobber each other / break MIC refs).
	// The key is an FName, so the collision check is case-insensitive — matching UObject name semantics.
	if (const FShaderLabModel* Existing = RegisteredModels.Find(ObjName))
	{
		if (!FPaths::IsSamePath(Existing->SourceFilePath, Model.SourceFilePath))
		{
			UE_LOG(LogShaderLabRegistry, Error,
				TEXT("ShaderLab: duplicate material name '%s' from '%s' and '%s' — rename one file."),
				*ObjName.ToString(), *Existing->SourceFilePath, *Model.SourceFilePath);
			return nullptr;
		}
	}

	UMaterial* Material = FindObject<UMaterial>(Pkg, *ObjName.ToString());
	const bool bNewlyCreated = (Material == nullptr);
	if (bNewlyCreated)
	{
		// RF_MarkAsRootSet keeps the material alive across GC. Without it the object is collected
		// (RF_Standalone alone is not a GC root in a cooked runtime), which removes its entry from
		// the IoStore script-object table (RemoveUnreachableScriptObjects) and makes cooked Material
		// Instances that reference it resolve their Parent to null -> fall back to the default material.
		Material = NewObject<UMaterial>(Pkg, ObjName, RF_Public | RF_Standalone | RF_MarkAsRootSet);
	}

	RegisteredModels.FindOrAdd(ObjName) = Model;

	// Tag the base material with the source `.usl` virtual path (stable across machines). MICs derive
	// their `ShaderLabPath` asset-registry tag from this; the editor resolves it back to disk (VSCode button).
	{
		UShaderLabMaterialAssetUserData* UserData = Material->GetAssetUserData<UShaderLabMaterialAssetUserData>();
		if (!UserData)
		{
			UserData = NewObject<UShaderLabMaterialAssetUserData>(Material);
			Material->AddAssetUserData(UserData);
		}
		UserData->ShaderLabPath = MakeVirtualShaderPath(Model.SourceFilePath);
	}

	// Register with the loading system so imports to "/Script/UShaderLab.<ObjName>" from other
	// packages resolve to this in-memory object — required for cooked builds, where the linker
	// resolves a Material Instance's parent import through this registration (without it the parent
	// resolves to null and the instance falls back to the default material).
	NotifyRegistrationEvent(
		FName(GShaderLabPackageName),
		ObjName,
		ENotifyRegistrationType::NRT_NoExportObject,
		ENotifyRegistrationPhase::NRP_Finished,
		nullptr,
		/*bDynamic*/ false,
		Material);

	BuildMaterialDelegate.Broadcast(*Material, Model);
	return Material;
}