UShaderLab/Source/UShaderLab/Private/ShaderLabMaterialInstanceConstant.cpp

// Copyright UShaderLab. All Rights Reserved.

#include "ShaderLabMaterialInstanceConstant.h"

#include "Engine/SpecularProfile.h"
#include "Engine/SubsurfaceProfile.h"
#include "Engine/ToonProfile.h"
#include "MaterialCachedData.h"
#include "Materials/Material.h"
#include "Materials/MaterialInstance.h"
#include "ShaderLabMaterialAssetUserData.h"
#include "ShaderLabMaterialRegistry.h"
#include "ShaderLabModel.h"
#include "UObject/AssetRegistryTagsContext.h"
#include "UObject/Package.h"

const FName UShaderLabMaterialInstanceConstant::ShaderLabPathTagName(TEXT("ShaderLabPath"));

const UMaterial* UShaderLabMaterialInstanceConstant::GetShaderLabRootBase() const
{
	// A ROOT instance's immediate parent is a non-instance UMaterial living in the ShaderLab /Script package.
	// Compare against the package POINTER (not a string name). Child instances (parent is another instance)
	// return nullptr — they defer to and share the root's shader map + derived data.
	if (const UMaterial* BaseMaterial = Cast<UMaterial>(Parent))
	{
		if (BaseMaterial->GetOutermost() == FShaderLabMaterialRegistry::Get().GetPackage())
		{
			return BaseMaterial;
		}
	}
	return nullptr;
}

bool UShaderLabMaterialInstanceConstant::HasOverridenBaseProperties() const
{
	// Force a self-contained permutation only at the ROOT — where our immediate parent is a ShaderLab
	// in-memory /Script base material (which has no serialized shader map). For child instances fall through
	// to stock behavior so they DEFER to and SHARE the root's shader map (bounds the shader-map count).
	if (GetShaderLabRootBase() != nullptr)
	{
		return true;
	}

	return Super::HasOverridenBaseProperties();
}

FString UShaderLabMaterialInstanceConstant::GetShaderLabPath(const UMaterialInterface* Material)
{
	// Walk parent -> ... -> ShaderLab base material; the base carries the source .usl virtual path.
	const UMaterialInterface* Current = Material;
	while (Current)
	{
		if (const UShaderLabMaterialAssetUserData* UserData =
			const_cast<UMaterialInterface*>(Current)->GetAssetUserData<UShaderLabMaterialAssetUserData>())
		{
			return UserData->ShaderLabPath;
		}
		const UMaterialInstance* Instance = Cast<UMaterialInstance>(Current);
		Current = Instance ? ToRawPtr(Instance->Parent) : nullptr;
	}
	return FString();
}

void UShaderLabMaterialInstanceConstant::PostLoad()
{
	Super::PostLoad();

#if WITH_EDITOR
	// The stock UMaterialInstanceConstant::UpdateCachedData() (called from Super::PostLoad) builds this
	// instance's OWN CachedExpressionData only when it has material layers; otherwise it leaves it null and
	// defers to the parent's. Our parent is the MEMORY-ONLY /Script base material, whose CachedExpressionData
	// is never serialized/cooked — so any /Game asset it references (Material Parameter Collections via
	// SL_COLLECTION, texture assets, functions) would be invisible to the cook's reference gathering AND never
	// loaded at runtime (e.g. an MPC's uniform buffer → fatal "Failed to find parameter collection buffer").
	//
	// Force the ROOT instance to build its OWN CachedExpressionData by analyzing the base graph. The engine's
	// standard UMaterialInterface::Serialize then harvests + serializes its hard references (see
	// MaterialInterface.cpp), so the cook includes them and the runtime loads them. The data is DERIVED from the
	// current base graph on every load/cook (never written to the editor .uasset), so editing the .usl needs no
	// MIC re-save — the memory base stays resource-transparent.
	BuildSelfContainedCachedExpressionData();
	ApplyShaderLabProfileOverrides();
#endif
}

#if WITH_EDITOR
void UShaderLabMaterialInstanceConstant::RefreshShaderLabDerivedData()
{
	if (GetShaderLabRootBase() == nullptr)
	{
		return; // Only root instances carry the derived cook-inclusion data.
	}

	// Invalidate the derived state so it rebuilds from the (just hot-reloaded) base graph. Without the reset,
	// BuildSelfContainedCachedExpressionData early-outs on the existing copy, and bLoadedCachedExpressionData=true
	// makes stock UpdateCachedData skip refreshing it — the resident MIC would keep stale references.
	CachedExpressionData.Reset();
	bLoadedCachedExpressionData = false;

	// Clear stale profile overrides too, so a profile removed from the .usl actually drops off (rather than
	// lingering because ApplyShaderLabProfileOverrides only ever sets, never clears).
	bOverrideSubsurfaceProfile = false;
	SubsurfaceProfile = nullptr;
	bOverrideSpecularProfile = false;
	SpecularProfileOverride = nullptr;
	bOverrideToonProfile = false;
	ToonProfileOverride = nullptr;

	BuildSelfContainedCachedExpressionData();
	ApplyShaderLabProfileOverrides();
}

void UShaderLabMaterialInstanceConstant::BuildSelfContainedCachedExpressionData()
{
	// Only the ROOT instance (immediate parent is a ShaderLab /Script base material) owns the self-contained
	// shader map; child instances share the root's map and its cached data (see HasOverridenBaseProperties).
	if (GetShaderLabRootBase() == nullptr)
	{
		return;
	}

	// Respect an existing build (e.g. a layered instance that already populated its own via the stock path).
	if (CachedExpressionData)
	{
		return;
	}

	UMaterial* BaseMaterial = GetMaterial();
	check(BaseMaterial); // A ShaderLab root instance always resolves to its /Script base material.

	// Ensure the base material's cached data reflects the current graph, then take our OWN copy of it. We copy
	// rather than call FMaterialCachedExpressionData::AnalyzeMaterial directly because that method is not
	// ENGINE_API-exported, whereas UpdateCachedExpressionData()/GetCachedExpressionData() are. The copy carries
	// the base graph's referenced parameter collections (with StateIds), texture assets and function infos.
	BaseMaterial->UpdateCachedExpressionData();
	CachedExpressionData = MakeUnique<FMaterialCachedExpressionData>(BaseMaterial->GetCachedExpressionData());

	// Mirror the editor-only side, exactly as the stock UMaterialInstanceConstant path does after a build.
	if (UMaterialInstanceEditorOnlyData* EditorData = GetEditorOnlyData())
	{
		EditorData->CachedExpressionData = CachedExpressionData->EditorOnlyData;
	}

	// CRITICAL: the stock UMaterialInstanceConstant::UpdateCachedData() rebuilds (for layered) or RESETS TO NULL
	// (for non-layered — our case) CachedExpressionData inside an `if (!bLoadedCachedExpressionData)` block. The
	// cook calls that again after PostLoad (via BeginCacheForCookedPlatformData → UpdateStaticPermutation), which
	// would wipe the copy we just made before it gets serialized. Marking it as "loaded" makes those subsequent
	// calls skip the reset, so our copy survives to cook-save (and is re-derived fresh on the next load/cook).
	bLoadedCachedExpressionData = true;
}

void UShaderLabMaterialInstanceConstant::ApplyShaderLabProfileOverrides()
{
	// Substrate profiles (Subsurface/Specular/Toon) are NOT part of FMaterialCachedExpressionData — they live in
	// UMaterial::SubsurfaceProfiles/SpecularProfiles/ToonProfiles on the memory-only /Script base and resolve via
	// the parent chain, so the CachedExpressionData copy above does not cover them and they would be missing from
	// the cook (→ profile-shaded demos render dark). Set them as this instance's OWN overrides (hard TObjectPtr
	// UPROPERTYs) so the cook harvests + serializes them and the runtime loads + applies them. Resolved from the
	// source `.usl` model (not the base's compiled Substrate info, which may not be populated yet at PostLoad).
	const UMaterial* ImmediateBase = GetShaderLabRootBase();
	if (!ImmediateBase)
	{
		return;
	}

	const TMap<FName, FShaderLabModel>& Models = FShaderLabMaterialRegistry::Get().GetRegisteredModels();
	const FShaderLabModel* Model = Models.Find(ImmediateBase->GetFName());
	if (!Model)
	{
		return;
	}

	// Collect the profile path of each type across the surface entry + all slabs. A MIC carries a SINGLE override
	// per profile type, and at runtime an instance override replaces the profile for ALL slots — so a shader using
	// two DIFFERENT profiles of the same type across slabs cannot be represented and would silently render both
	// with one profile. Fail loud on that (contract-style) rather than ship a wrong material; repeats of the SAME
	// path are fine.
	FString SubsurfacePath, SpecularPath, ToonPath;
	auto Collect = [Model](FString& Dst, const FString& Src, const TCHAR* Type)
	{
		if (Src.IsEmpty())
		{
			return;
		}
		checkf(Dst.IsEmpty() || Dst == Src,
			TEXT("ShaderLab '%s': two distinct %s profiles across slabs ('%s' vs '%s'); a material instance carries ")
			TEXT("one override per profile type, so multi-profile-per-type is unsupported. Use a single %s profile."),
			*Model->ShaderName, Type, *Dst, *Src, Type);
		Dst = Src;
	};
	Collect(SubsurfacePath, Model->SurfaceModifiers.SubsurfaceProfilePath, TEXT("Subsurface"));
	Collect(SpecularPath, Model->SurfaceModifiers.SpecularProfilePath, TEXT("Specular"));
	Collect(ToonPath, Model->SurfaceModifiers.ToonProfilePath, TEXT("Toon"));
	for (const FShaderLabSlab& Slab : Model->Slabs)
	{
		Collect(SubsurfacePath, Slab.Modifiers.SubsurfaceProfilePath, TEXT("Subsurface"));
		Collect(SpecularPath, Slab.Modifiers.SpecularProfilePath, TEXT("Specular"));
		Collect(ToonPath, Slab.Modifiers.ToonProfilePath, TEXT("Toon"));
	}

	// Contract-style: an author-declared profile path that fails to load is a hard error (it would silently ship a
	// dark/incorrect material) — the same class of bug this override exists to prevent.
	if (!SubsurfacePath.IsEmpty())
	{
		USubsurfaceProfile* Profile = LoadObject<USubsurfaceProfile>(nullptr, *SubsurfacePath);
		checkf(Profile, TEXT("ShaderLab '%s': SubsurfaceProfile '%s' failed to load."), *Model->ShaderName, *SubsurfacePath);
		bOverrideSubsurfaceProfile = true;
		SubsurfaceProfile = Profile;
	}
	if (!SpecularPath.IsEmpty())
	{
		USpecularProfile* Profile = LoadObject<USpecularProfile>(nullptr, *SpecularPath);
		checkf(Profile, TEXT("ShaderLab '%s': SpecularProfile '%s' failed to load."), *Model->ShaderName, *SpecularPath);
		bOverrideSpecularProfile = true;
		SpecularProfileOverride = Profile;
	}
	if (!ToonPath.IsEmpty())
	{
		UToonProfile* Profile = LoadObject<UToonProfile>(nullptr, *ToonPath);
		checkf(Profile, TEXT("ShaderLab '%s': ToonProfile '%s' failed to load."), *Model->ShaderName, *ToonPath);
		bOverrideToonProfile = true;
		ToonProfileOverride = Profile;
	}
}

void UShaderLabMaterialInstanceConstant::GetAssetRegistryTags(FAssetRegistryTagsContext Context) const
{
	Super::GetAssetRegistryTags(Context);

	// Export the source .usl virtual path so users can filter ShaderLab instances in the Content Browser.
	const FString Path = GetShaderLabPath(this);
	if (!Path.IsEmpty())
	{
		Context.AddTag(FAssetRegistryTag(ShaderLabPathTagName, Path, FAssetRegistryTag::TT_Alphabetical));
	}
}
#endif