Files
UShaderLab/Source/UShaderLabEditor/Private/ShaderLabEditorModule.cpp
2026-07-03 17:44:09 +08:00

339 lines
12 KiB
C++

// Copyright UShaderLab. All Rights Reserved.
#include "DirectoryWatcherModule.h"
#include "Editor.h"
#include "Engine/Engine.h"
#include "HAL/FileManager.h"
#include "IDirectoryWatcher.h"
#include "Interfaces/IPluginManager.h"
#include "MaterialEditingLibrary.h"
#include "MaterialShared.h"
#include "Materials/Material.h"
#include "Misc/Paths.h"
#include "Modules/ModuleManager.h"
#include "ShaderCore.h"
#include "Misc/FileHelper.h"
#include "ShaderLabDiscovery.h"
#include "ShaderLabGraphBuilder.h"
#include "ShaderLabImportResolver.h"
#include "ShaderLabMaterialInstanceConstant.h"
#include "ShaderLabMaterialRegistry.h"
#include "ShaderLabModel.h"
#include "ShaderLabParser.h"
#include "ShaderLabSubsystem.h"
#include "ShaderLabVSCodeButton.h"
#include "UObject/Package.h"
#include "UObject/UObjectIterator.h"
DEFINE_LOG_CATEGORY_STATIC(LogShaderLabEditor, Log, All);
namespace
{
/** Editor-side reaction to a base material needing build: rebuild graph + trigger compile. */
void BuildAndCompile(UMaterial& Material, const FShaderLabModel& Model)
{
// Absolute, forward-slashed source path for the poison node's `#line` mapping (matches the graph
// builder's own convention so compiler errors click through to the .usl).
FString SrcPath = FPaths::ConvertRelativePathToFull(Model.SourceFilePath);
SrcPath.ReplaceInline(TEXT("\\"), TEXT("/"));
if (Model.LoadErrors.Num() > 0)
{
// Parse-stage failure: the .usl never produced a valid model. Route the diagnostics through
// the shader compiler as a poison material so they surface in the MIC editor / cook, not just
// the log.
for (const FString& E : Model.LoadErrors)
{
UE_LOG(LogShaderLabEditor, Error, TEXT("ShaderLab parse '%s': %s"), *Model.ShaderName, *E);
}
FShaderLabGraphBuilder::BuildPoisonInto(Material, Model.LoadErrors, SrcPath);
}
else
{
TArray<FString> Errors;
if (!FShaderLabGraphBuilder::BuildInto(Material, Model, Errors))
{
for (const FString& E : Errors)
{
UE_LOG(LogShaderLabEditor, Error, TEXT("ShaderLab build '%s': %s"), *Model.ShaderName, *E);
}
// BuildInto already cleared the graph before failing — don't leave the base silently empty.
// Replace it with a poison material so the build errors compile-fail with the same
// visibility as an HLSL error, instead of vanishing into the log.
FShaderLabGraphBuilder::BuildPoisonInto(Material, Errors, SrcPath);
}
}
// Recompile the base AND propagate to everything that depends on it. A bare PostEditChange
// recompiles only the base shell; the instances carry their own static-permutation shader maps
// (and the level components render those), so on a hot-reload they must be recompiled and
// re-pushed too. FMaterialUpdateContext does exactly that on scope exit: recompile the listed
// material, then recache every dependent instance and refresh the components/viewports using it.
{
FMaterialUpdateContext UpdateContext;
UpdateContext.AddMaterial(&Material);
Material.PreEditChange(nullptr);
Material.PostEditChange();
}
// Re-derive the cook-inclusion data (CachedExpressionData copy + profile overrides) on every ROOT
// UShaderLabMaterialInstanceConstant of this base. The FMaterialUpdateContext above recompiles the
// dependent shader maps but does NOT refresh that derived data (it's built once at PostLoad and pinned by
// bLoadedCachedExpressionData=true), so a hot-reload that added/removed an MPC or profile would otherwise
// leave a resident MIC stale — and an in-editor iterative cook would serialize the stale references.
//
// Inheritance chain: only a ROOT (immediate parent == this base) carries the derived data; CHILD instances
// (parent is another instance, however deep) hold none and resolve up the chain to the root at query time,
// and their shared shader maps were just recompiled by FMaterialUpdateContext — so refreshing the roots is
// sufficient and children need no separate pass. RefreshShaderLabDerivedData self-guards via
// GetShaderLabRootBase(), so it is a safe no-op on anything that isn't a root of this base. Skip CDOs and
// objects being GC'd.
for (TObjectIterator<UShaderLabMaterialInstanceConstant> It(
/*AdditionalExclusionFlags*/ RF_ClassDefaultObject,
/*bIncludeDerivedClasses*/ true,
/*InternalExclusionFlags*/ EInternalObjectFlags::Garbage); It; ++It)
{
if (It->Parent == &Material)
{
It->RefreshShaderLabDerivedData();
}
}
// Refresh any open Material Instance editors built on this base so newly added/removed parameters
// (scalars, textures, static-bool switches) appear without reopening the editor. RegenerateArrays
// pulls the parameter set from the freshly rebuilt parent; mirrors what the stock material editor
// does after RecompileMaterial. Must run after PostEditChange so the parent's parameters are current.
// Guarded on GEditor: this handler also runs during cook/headless registration, where there is no
// interactive editor (and RebuildMaterialInstanceEditors dereferences GEditor).
if (GEditor)
{
UMaterialEditingLibrary::RebuildMaterialInstanceEditors(&Material);
}
}
}
class FShaderLabEditorModule : public IModuleInterface
{
public:
virtual void StartupModule() override
{
// Map the plugin's Shaders/ directory so generated Custom nodes can #include
// "/Plugin/ShaderLab/Private/ShaderLabCommon.ush" at shader-compile time.
if (const TSharedPtr<IPlugin> Plugin = IPluginManager::Get().FindPlugin(TEXT("UShaderLab")))
{
const FString ShaderDir = FPaths::Combine(Plugin->GetBaseDir(), TEXT("Shaders"));
if (!AllShaderSourceDirectoryMappings().Contains(TEXT("/Plugin/ShaderLab")))
{
AddShaderSourceDirectoryMapping(TEXT("/Plugin/ShaderLab"), ShaderDir);
}
}
// Map the project's Shaders/ directory to the virtual root "/Project" so user .ush library
// files placed there (alongside .usl) can be #included from a shader's Includes { } block.
{
const FString ProjectShaderDir = FPaths::Combine(FPaths::ProjectDir(), TEXT("Shaders"));
if (FPaths::DirectoryExists(ProjectShaderDir) && !AllShaderSourceDirectoryMappings().Contains(TEXT("/Project")))
{
AddShaderSourceDirectoryMapping(TEXT("/Project"), ProjectShaderDir);
}
}
// Map the plugin's Intermediate/ShaderLabGen directory (a build-time artifact dir) so the graph
// builder can #include per-shader generated local-code headers (`/UShaderLabGen/<Name>.gen.ush`).
// AddShaderSourceDirectoryMapping requires the real dir to already exist, so create it first.
{
const FString GenDir = FShaderLabGraphBuilder::GetGeneratedShaderDir();
const TCHAR* GenRoot = FShaderLabGraphBuilder::GetGeneratedVirtualRoot();
if (!GenDir.IsEmpty())
{
IFileManager::Get().MakeDirectory(*GenDir, /*Tree*/ true);
if (FPaths::DirectoryExists(GenDir) && !AllShaderSourceDirectoryMappings().Contains(GenRoot))
{
AddShaderSourceDirectoryMapping(GenRoot, GenDir);
}
}
}
// Fill in / recompile base materials whenever the registry asks (startup + hot reload).
BuildHandle = FShaderLabMaterialRegistry::Get().OnBuildMaterial().AddStatic(&BuildAndCompile);
StartWatchingSources();
// "Open .usl in VSCode" toolbar button on the Material Instance Editor for ShaderLab MICs.
ShaderLabVSCodeButton::Register();
UE_LOG(LogShaderLabEditor, Log, TEXT("ShaderLabEditor module started."));
}
virtual void ShutdownModule() override
{
StopWatchingSources();
ShaderLabVSCodeButton::Unregister();
if (BuildHandle.IsValid())
{
FShaderLabMaterialRegistry::Get().OnBuildMaterial().Remove(BuildHandle);
BuildHandle.Reset();
}
}
private:
void StartWatchingSources()
{
FDirectoryWatcherModule& Module =
FModuleManager::LoadModuleChecked<FDirectoryWatcherModule>(TEXT("DirectoryWatcher"));
IDirectoryWatcher* Watcher = Module.Get();
if (!Watcher)
{
return;
}
for (const FString& Root : FShaderLabDiscovery::GetSearchRoots())
{
if (!FPaths::DirectoryExists(Root))
{
continue;
}
FDelegateHandle Handle;
Watcher->RegisterDirectoryChangedCallback_Handle(
Root,
IDirectoryWatcher::FDirectoryChanged::CreateStatic(&FShaderLabEditorModule::OnDirectoryChanged),
Handle,
IDirectoryWatcher::WatchOptions::IncludeDirectoryChanges);
WatchedRoots.Add(Root, Handle);
}
}
void StopWatchingSources()
{
if (FModuleManager::Get().IsModuleLoaded(TEXT("DirectoryWatcher")))
{
FDirectoryWatcherModule& Module =
FModuleManager::GetModuleChecked<FDirectoryWatcherModule>(TEXT("DirectoryWatcher"));
if (IDirectoryWatcher* Watcher = Module.Get())
{
for (const TPair<FString, FDelegateHandle>& Pair : WatchedRoots)
{
Watcher->UnregisterDirectoryChangedCallback_Handle(Pair.Key, Pair.Value);
}
}
}
WatchedRoots.Reset();
}
/** Resolve a virtual shader path to disk via the registered shader-source directory mappings. */
static bool ResolveVirtualToDisk(const FString& VirtualPath, FString& OutDiskPath)
{
FString Best, BestDir;
for (const TPair<FString, FString>& Pair : AllShaderSourceDirectoryMappings())
{
if ((VirtualPath.StartsWith(Pair.Key + TEXT("/")) || VirtualPath == Pair.Key) && Pair.Key.Len() > Best.Len())
{
Best = Pair.Key;
BestDir = Pair.Value;
}
}
if (Best.IsEmpty()) { return false; }
FString Rest = VirtualPath.Mid(Best.Len());
Rest.RemoveFromStart(TEXT("/"));
OutDiskPath = FPaths::Combine(BestDir, Rest);
return true;
}
static FString NormalizePath(const FString& Path)
{
FString N = FPaths::ConvertRelativePathToFull(Path);
N.ReplaceInline(TEXT("\\"), TEXT("/"));
return N;
}
/**
* Every discovered `.usl` shader whose transitive `.uslfunc` import closure contains ChangedDiskPath —
* i.e. the shaders that must be rebuilt when that library changes. Resolves each shader's imports so
* both direct and nested (library-imports-library) dependencies are caught.
*/
static TArray<FString> FindShadersImporting(const FString& ChangedDiskPathNormalized)
{
TArray<FString> Result;
for (const FString& Usl : FShaderLabDiscovery::FindShaderLabFiles())
{
FString Source;
if (!FFileHelper::LoadFileToString(Source, *Usl))
{
continue;
}
FShaderLabModel Model;
TArray<FShaderLabParseError> ParseErrors;
if (!FShaderLabParser::Parse(Source, Usl, Model, ParseErrors) || Model.Imports.Num() == 0)
{
continue;
}
FShaderLabImportResolver::FSourceLoader Loader =
[](const FString& VPath, FString& OutSrc, FString& OutDisk, FString& OutErr) -> bool
{
if (!ResolveVirtualToDisk(VPath, OutDisk)) { OutErr = TEXT("unmapped"); return false; }
if (!FFileHelper::LoadFileToString(OutSrc, *OutDisk)) { OutErr = TEXT("read failed"); return false; }
return true;
};
FShaderLabResolvedProgram Program;
TArray<FString> ResolveErrors;
if (!FShaderLabImportResolver::Resolve(Model, Loader, Program, ResolveErrors))
{
continue;
}
for (const FShaderLabResolvedLibrary& Lib : Program.Libraries)
{
if (NormalizePath(Lib.Model.SourceFilePath) == ChangedDiskPathNormalized)
{
Result.Add(Usl);
break;
}
}
}
return Result;
}
static void OnDirectoryChanged(const TArray<FFileChangeData>& Changes)
{
UShaderLabSubsystem* Subsystem = GEngine ? GEngine->GetEngineSubsystem<UShaderLabSubsystem>() : nullptr;
if (!Subsystem)
{
return;
}
TSet<FString> Rebuilt;
auto RebuildOnce = [&](const FString& Normalized)
{
if (Rebuilt.Contains(Normalized)) { return; }
Rebuilt.Add(Normalized);
Subsystem->RebuildFromFile(Normalized);
UE_LOG(LogShaderLabEditor, Log, TEXT("ShaderLab hot-reloaded '%s'"), *Normalized);
};
for (const FFileChangeData& Change : Changes)
{
if (Change.Action == FFileChangeData::FCA_Removed)
{
continue; // Leave the existing in-memory material in place on delete.
}
const FString Ext = FPaths::GetExtension(Change.Filename);
const FString Normalized = NormalizePath(Change.Filename);
if (Ext.Equals(TEXT("usl"), ESearchCase::IgnoreCase))
{
RebuildOnce(Normalized);
}
else if (Ext.Equals(TEXT("uslfunc"), ESearchCase::IgnoreCase))
{
// A library changed: rebuild every shader that (transitively) imports it, so their promoted
// parameters and emitted function HLSL refresh — the dependency-aware hot reload.
for (const FString& Dependent : FindShadersImporting(Normalized))
{
RebuildOnce(NormalizePath(Dependent));
}
}
}
}
FDelegateHandle BuildHandle;
TMap<FString, FDelegateHandle> WatchedRoots;
};
IMPLEMENT_MODULE(FShaderLabEditorModule, UShaderLabEditor)