UShaderLab/Source/UShaderLabEditor/Private/ShaderLabEditorModule.cpp

// Copyright FlecsProj. All Rights Reserved.

#include "DirectoryWatcherModule.h"
#include "Engine/Engine.h"
#include "IDirectoryWatcher.h"
#include "Interfaces/IPluginManager.h"
#include "Materials/Material.h"
#include "Misc/Paths.h"
#include "Modules/ModuleManager.h"
#include "ShaderCore.h"
#include "ShaderLabDiscovery.h"
#include "ShaderLabGraphBuilder.h"
#include "ShaderLabMaterialRegistry.h"
#include "ShaderLabModel.h"
#include "ShaderLabSubsystem.h"
#include "UObject/Package.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)
	{
		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);
			}
			return;
		}
		// Recompile for rendering / editor preview (async).
		Material.PostEditChange();
	}
}

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);
			}
		}

		// Fill in / recompile base materials whenever the registry asks (startup + hot reload).
		BuildHandle = FShaderLabMaterialRegistry::Get().OnBuildMaterial().AddStatic(&BuildAndCompile);

		StartWatchingSources();

		// Note: ShaderLab base materials are deliberately NOT surfaced to the Content Browser — they are
		// hidden templates. Users create UShaderLabMaterialInstanceConstant instances via the factory
		// (Content Browser ▸ Material ▸ ShaderLab Material Instance), which picks a base directly from
		// the registry. The instances are the only assets users interact with.
		//
		// There is NO cook-time hook: the .usl sources are staged into packaged builds (a wildcard
		// RuntimeDependency in ShaderLab.Build.cs) and the cooked runtime re-parses them directly, so
		// nothing special needs to be written at cook.

		UE_LOG(LogShaderLabEditor, Log, TEXT("ShaderLabEditor module started."));
	}

	virtual void ShutdownModule() override
	{
		StopWatchingSources();
		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();
	}

	static void OnDirectoryChanged(const TArray<FFileChangeData>& Changes)
	{
		UShaderLabSubsystem* Subsystem = GEngine ? GEngine->GetEngineSubsystem<UShaderLabSubsystem>() : nullptr;
		if (!Subsystem)
		{
			return;
		}
		TSet<FString> Rebuilt;
		for (const FFileChangeData& Change : Changes)
		{
			if (!Change.Filename.EndsWith(TEXT(".usl")))
			{
				continue;
			}
			if (Change.Action == FFileChangeData::FCA_Removed)
			{
				continue; // Leave the existing in-memory material in place on delete.
			}
			const FString Normalized = FPaths::ConvertRelativePathToFull(Change.Filename);
			if (Rebuilt.Contains(Normalized))
			{
				continue;
			}
			Rebuilt.Add(Normalized);
			Subsystem->RebuildFromFile(Normalized);
			UE_LOG(LogShaderLabEditor, Log, TEXT("ShaderLab hot-reloaded '%s'"), *Normalized);
		}
	}

	FDelegateHandle BuildHandle;
	TMap<FString, FDelegateHandle> WatchedRoots;
};

IMPLEMENT_MODULE(FShaderLabEditorModule, UShaderLabEditor)