UnrealEngineUWP/Engine/Source/Programs/Shared/EpicGames.Build/System/Rules.cs

// Copyright Epic Games, Inc. All Rights Reserved.

using EpicGames.Core;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;

namespace UnrealBuildBase
{
	public class Rules
	{
		/// <summary>
		/// Enum for types of rules files. Should match extensions in RulesFileExtensions.
		/// </summary>
		public enum RulesFileType
		{
			/// <summary>
			/// .build.cs files
			/// </summary>
			Module,

			/// <summary>
			/// .target.cs files
			/// </summary>
			Target,

			/// <summary>
			/// .automation.csproj files
			/// </summary>
			AutomationModule,

			/// <summary>
			/// .ubtplugin.csproj files
			/// </summary>
			UbtPlugin,
		}

		/// <summary>
		/// Cached list of rules files in each directory of each type
		/// </summary>
		class RulesFileCache
		{
			public List<FileReference> ModuleRules = new List<FileReference>();
			public List<FileReference> TargetRules = new List<FileReference>();
			public List<FileReference> AutomationModules = new List<FileReference>();
			public List<FileReference> UbtPlugins = new List<FileReference>();
		}

		/// Map of root folders to a cached list of all UBT-related source files in that folder or any of its sub-folders.
		/// We cache these file names so we can avoid searching for them later on.
		static Dictionary<DirectoryReference, RulesFileCache> RootFolderToRulesFileCache = new Dictionary<DirectoryReference, RulesFileCache>();

		/// <summary>
		/// 
		/// </summary>
		/// <param name="RulesFileType"></param>
		/// <param name="GameFolders"></param>
		/// <param name="ForeignPlugins"></param>
		/// <param name="AdditionalSearchPaths"></param>
		/// <param name="bIncludeEngine"></param>
		/// <param name="bIncludeTempTargets">Whether to include targets generated by UAT to accomodate content-only projects that need to be compiled to include plugins</param>
		/// <returns></returns>
		public static List<FileReference> FindAllRulesSourceFiles(RulesFileType RulesFileType, List<DirectoryReference>? GameFolders, List<FileReference>? ForeignPlugins, List<DirectoryReference>? AdditionalSearchPaths, bool bIncludeEngine = true, bool bIncludeTempTargets = true)
		{
			List<DirectoryReference> Folders = new List<DirectoryReference>();

			// Add all engine source (including third party source)
			if (bIncludeEngine)
			{
				Folders.AddRange(Unreal.GetExtensionDirs(Unreal.EngineDirectory, "Source"));
			}

			// @todo plugin: Disallow modules from including plugin modules as dependency modules? (except when the module is part of that plugin)

			// Get all the root folders for plugins
			List<DirectoryReference> RootFolders = new List<DirectoryReference>();
			if (bIncludeEngine)
			{
				RootFolders.AddRange(Unreal.GetExtensionDirs(Unreal.EngineDirectory));
			}
			if (GameFolders != null)
			{
				RootFolders.AddRange(GameFolders.SelectMany(x => Unreal.GetExtensionDirs(x)));
			}

			// Find all the plugin source and tests directories
			foreach (DirectoryReference RootFolder in RootFolders)
			{
				DirectoryReference PluginsFolder = DirectoryReference.Combine(RootFolder, "Plugins");
				foreach (FileReference PluginFile in PluginsBase.EnumeratePlugins(PluginsFolder))
				{
					Folders.Add(DirectoryReference.Combine(PluginFile.Directory, "Source"));
					Folders.Add(DirectoryReference.Combine(PluginFile.Directory, "Tests"));
				}
			}

			// Add all the extra plugin folders
			if (ForeignPlugins != null)
			{
				foreach (FileReference ForeignPlugin in ForeignPlugins)
				{
					Folders.Add(DirectoryReference.Combine(ForeignPlugin.Directory, "Source"));
				}
			}

			// Add in the game folders to search
			if (GameFolders != null)
			{
				foreach (DirectoryReference GameFolder in GameFolders)
				{
					Folders.AddRange(Unreal.GetExtensionDirs(GameFolder, "Source"));

					if (bIncludeTempTargets)
					{
						DirectoryReference GameIntermediateSourceFolder = DirectoryReference.Combine(GameFolder, "Intermediate", "Source");
						Folders.Add(GameIntermediateSourceFolder);
					}
				}
			}

			// Process the additional search path, if sent in, only for paths that exist
			if (AdditionalSearchPaths != null)
			{
				Folders.AddRange(AdditionalSearchPaths.Where(x => x != null && DirectoryReference.Exists(x)));
			}

			return FindAllRulesFiles(Folders, RulesFileType);
		}

		/// <summary>
		/// Invalidate the cache for the givcen directory
		/// </summary>
		/// <param name="DirectoryPath">Directory to invalidate</param>
        public static void InvalidateRulesFileCache(string DirectoryPath)
        {
            DirectoryReference Directory = new DirectoryReference(DirectoryPath);
			lock(RootFolderToRulesFileCache)
			{
				RootFolderToRulesFileCache.Remove(Directory);
			}
            DirectoryLookupCache.InvalidateCachedDirectory(Directory);
        }

		/// <summary>
		/// Prefetch multiple directories in parallel
		/// </summary>
		/// <param name="Directories">The directories to cache</param>
		public static void PrefetchRulesFiles(IEnumerable<DirectoryReference> Directories)
		{
			PrefetchRulesFilesInternal(Directories);
		}

		/// <summary>
		/// Prefetch multiple directories in parallel
		/// </summary>
		/// <param name="Directories">The directories to cache</param>
		/// <returns>The list of caches for each directory</returns>
		private static List<(DirectoryReference, RulesFileCache)> PrefetchRulesFilesInternal(IEnumerable<DirectoryReference> Directories)
		{
			List<(DirectoryReference, RulesFileCache)> Caches = new List<(DirectoryReference, RulesFileCache)>(Directories.Count());
			lock (RootFolderToRulesFileCache)
			{
				using (ThreadPoolWorkQueue Queue = new ThreadPoolWorkQueue())
				{
					foreach (DirectoryReference Directory in Directories)
					{
						RulesFileCache? Cache;
						if (!RootFolderToRulesFileCache.TryGetValue(Directory, out Cache))
						{
							Cache = new RulesFileCache();
							Queue.Enqueue(() => FindAllRulesFilesRecursively(DirectoryItem.GetItemByDirectoryReference(Directory), Cache, Queue));
						}
						Caches.Add((Directory, Cache));
					}
				}

				using (ThreadPoolWorkQueue SortQueue = new ThreadPoolWorkQueue())
				{
					foreach ((DirectoryReference Directory, RulesFileCache Cache) in Caches)
					{
						if (!RootFolderToRulesFileCache.ContainsKey(Directory))
						{
							SortQueue.Enqueue(() =>
							{
								Cache.ModuleRules.Sort((A, B) => A.FullName.CompareTo(B.FullName));
								Cache.TargetRules.Sort((A, B) => A.FullName.CompareTo(B.FullName));
								Cache.AutomationModules.Sort((A, B) => A.FullName.CompareTo(B.FullName));
								Cache.UbtPlugins.Sort((A, B) => A.FullName.CompareTo(B.FullName));
							});
							RootFolderToRulesFileCache[Directory] = Cache;
						}
					}
				}
			}
			return Caches;
		}

		/// <summary>
		/// Finds all the rules of the given type under a given directory
		/// </summary>
		/// <param name="Directory">Directory to search</param>
		/// <param name="Type">Type of rules to return</param>
		/// <returns>List of rules files of the given type</returns>
		public static IReadOnlyList<FileReference> FindAllRulesFiles(DirectoryReference Directory, RulesFileType Type)
		{
			return FindAllRulesFiles(new List<DirectoryReference> { Directory }, Type);
		}

		/// <summary>
		/// Finds all the rules of the given type under a given directories
		/// </summary>
		/// <param name="Directories">Directories to search</param>
		/// <param name="Type">Type of rules to return</param>
		/// <returns>List of rules files of the given type</returns>
		public static List<FileReference> FindAllRulesFiles(IEnumerable<DirectoryReference> Directories, RulesFileType Type)
		{
			List<(DirectoryReference, RulesFileCache)> Caches = PrefetchRulesFilesInternal(Directories);
			List<FileReference> Files = new List<FileReference>();			
			foreach ((DirectoryReference Directory, RulesFileCache Cache) in Caches)
			{
				// Get the list of files of the type we're looking for
				if (Type == RulesFileType.Module)
				{
					Files.AddRange(Cache.ModuleRules);
				}
				else if (Type == RulesFileType.Target)
				{
					Files.AddRange(Cache.TargetRules);
				}
				else if (Type == RulesFileType.AutomationModule)
				{
					Files.AddRange(Cache.AutomationModules);
				}
				else if (Type == RulesFileType.UbtPlugin)
				{
					Files.AddRange(Cache.UbtPlugins);
				}
				else
				{
					throw new Exception($"Unhandled rules type: {Type}");
				}
			}
			return Files;
		}

		/// <summary>
		/// Search through a directory tree for any rules files
		/// </summary>
		/// <param name="Directory">The root directory to search from</param>
		/// <param name="Cache">Receives all the discovered rules files</param>
		/// <param name="Queue">Queue for adding additional tasks to</param>
		private static void FindAllRulesFilesRecursively(DirectoryItem Directory, RulesFileCache Cache, ThreadPoolWorkQueue Queue)
		{
			// Scan all the files in this directory
			bool bSearchSubFolders = true;
			bool bIsPlugin = false;

			foreach (FileItem File in Directory.EnumerateFiles())
			{
				if (File.HasExtension(".build.cs"))
				{
					lock(Cache.ModuleRules)
					{
						Cache.ModuleRules.Add(File.Location);
					}
					bSearchSubFolders = false;
				}
				else if (File.HasExtension(".target.cs"))
				{
					lock(Cache.TargetRules)
					{
						Cache.TargetRules.Add(File.Location);
					}
				}
				else if (File.HasExtension(".automation.csproj"))
				{
					lock(Cache.AutomationModules)
					{
						Cache.AutomationModules.Add(File.Location);
					}
					bSearchSubFolders = false;
				}
				else if (File.HasExtension(".ubtplugin.csproj"))
				{
					lock(Cache.UbtPlugins)
					{
						Cache.UbtPlugins.Add(File.Location);
					}
					bSearchSubFolders = false;
				}
				else if (File.HasExtension(".uplugin"))
				{
					bIsPlugin = true;
					bSearchSubFolders = false;
				}
				else if (File.Name == ".ubtignore")
				{
					bSearchSubFolders = false;
				}
			}

			if (bIsPlugin)
			{
				DirectoryItem? SourceDir;
				if (Directory.TryGetDirectory("Source", out SourceDir))
				{
					Queue.Enqueue(() => FindAllRulesFilesRecursively(SourceDir, Cache, Queue));
				}
			}

			// If we didn't find anything to stop the search, search all the subdirectories too
			if (bSearchSubFolders)
			{
				foreach (DirectoryItem SubDirectory in Directory.EnumerateDirectories())
				{
					Queue.Enqueue(() => FindAllRulesFilesRecursively(SubDirectory, Cache, Queue));
				}
			}
		}
	}
}