UShaderLab/Source/UShaderLabEditor/Private/ShaderLabEditorModule.cpp

// Copyright UShaderLab. All Rights Reserved.

#include "DirectoryWatcherModule.h"
#include "Editor.h"
#include "Engine/Engine.h"
#include "IDirectoryWatcher.h"
#include "MaterialEditingLibrary.h"
#include "Materials/Material.h"
#include "Misc/FileHelper.h"
#include "Misc/Paths.h"
#include "Modules/ModuleManager.h"
#include "ShaderLabDiscovery.h"
#include "ShaderLabGraphBuilder.h"
#include "ShaderLabImportResolver.h"
#include "ShaderLabMaterialRegistry.h"
#include "ShaderLabModel.h"
#include "ShaderLabParser.h"
#include "ShaderLabSubsystem.h"
#include "ShaderLabVSCodeButton.h"

DEFINE_LOG_CATEGORY_STATIC(LogShaderLabEditor, Log, All);

// This module hosts only the interactive-editor UX for ShaderLab (GIsEditor-gated): the .usl directory
// watcher for hot reload, the "open in VSCode" toolbar button, the Material Instance factory, and
// refreshing open Material Instance editors after a rebuild. The material-graph-building core moved to
// UShaderLabBuilder (an UncookedOnly module) so it also runs under `-game`, where this Editor module and
// its editor-UI dependencies (UnrealEd/MaterialEditor) are absent.

namespace
{
	/**
	 * Editor-only reaction to a base material finishing its (re)build: refresh any open Material Instance
	 * editors 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. Bound to OnMaterialGraphBuilt (not
	 * OnBuildMaterial) so it runs strictly AFTER the builder finished the graph + compile. Guarded on GEditor:
	 * the builder also runs during cook/headless registration and under `-game`, where there is no interactive
	 * editor (and RebuildMaterialInstanceEditors dereferences GEditor).
	 */
	void RefreshInstanceEditors(UMaterial& Material)
	{
		if (GEditor)
		{
			UMaterialEditingLibrary::RebuildMaterialInstanceEditors(&Material);
		}
	}
}

class FShaderLabEditorModule : public IModuleInterface
{
public:
	virtual void StartupModule() override
	{
		// After the builder rebuilds a base material's graph, refresh any open MIC editors on it.
		GraphBuiltHandle = FShaderLabMaterialRegistry::Get().OnMaterialGraphBuilt().AddStatic(&RefreshInstanceEditors);

		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 (GraphBuiltHandle.IsValid())
		{
			FShaderLabMaterialRegistry::Get().OnMaterialGraphBuilt().Remove(GraphBuiltHandle);
			GraphBuiltHandle.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();
	}

	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 (!FShaderLabGraphBuilder::ResolveVirtualShaderFile(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 GraphBuiltHandle;
	TMap<FString, FDelegateHandle> WatchedRoots;
};

IMPLEMENT_MODULE(FShaderLabEditorModule, UShaderLabEditor)