NRDUnityIntegration.cs

namespace nrd
{
    public static class NRDUnityIntegration
    {

        [MenuItem("Rendering/Generate NRD Unity Code", false, 2500)]
        static void GenerateNRDCode()
        {
            GenerateUnityComputeShaders();
            GenerateCSharpShaderBindings();
        }

        static void GenerateUnityComputeShaders()
        {
            DirectoryInfo directoryInfo = new DirectoryInfo(nrdShaderFolder + "Source/");
            FileInfo[] hlslFiles = directoryInfo.GetFiles("*.cs.hlsl");
            foreach (FileInfo file in hlslFiles)
            {
                StreamWriter writer = File.CreateText(file.FullName.Replace(".cs.hlsl", ".compute"));

                writer.WriteLine("#pragma kernel main");
                writer.WriteLine(nrdCompilerConfiguration);
                writer.WriteLine(string.Format("#include \"../Include/NRDEncoding.hlsli\""));
                writer.WriteLine(string.Format("#include \"{0}\"", file.Name));

                writer.Flush();
                writer.Close();
            }
        }

        static void GenerateCSharpShaderBindings([CallerFilePath] string sourceFilePath = "")
        {
            StreamWriter writer = File.CreateText(Path.Combine(Path.GetDirectoryName(sourceFilePath), "NRDShaderBindings.cs"));

            writer.WriteLine("using float4x4 = UnityEngine.Matrix4x4;");
            writer.WriteLine("using float4 = UnityEngine.Vector4;");
            writer.WriteLine("using float3 = UnityEngine.Vector3;");
            writer.WriteLine("using float2 = UnityEngine.Vector2;");
            writer.WriteLine("using uint2 = UnityEngine.Vector2Int;");
            writer.WriteLine("using int2 = UnityEngine.Vector2Int;");
            writer.WriteLine();

            writer.WriteLine("namespace nrd {");
            writer.WriteLine();

            writer.WriteLine(reblurSharedCBData);
            writer.WriteLine(relaxSharedCBData);
            writer.WriteLine(sigmaSharedCBData);

            DirectoryInfo directoryInfo = new DirectoryInfo(nrdShaderFolder + "Resources/");
            FileInfo[] hlslFiles = directoryInfo.GetFiles("*.resources.hlsli");
            foreach (FileInfo file in hlslFiles)
            {
                writer.WriteLine(string.Format("// Generated from {0}", file.Name));

                string text = file.OpenText().ReadToEnd();

                foreach (string def in RegexMatchesStartEndMultiline(text, "NRD_CONSTANTS_START", "NRD_CONSTANTS_END"))
                {
                    string structName = file.Name.Replace(".resources.hlsli", "_Constants");
                    string structDef = def.Replace("NRD_CONSTANTS_START", "public struct " + structName + " {")
                                          .Replace("NRD_CONSTANTS_END", "}")
                                          .Replace("REBLUR_SHARED_CB_DATA", "public REBLUR_SHARED_CB_DATA Common;")
                                          .Replace("RELAX_SHARED_CB_DATA", "public RELAX_SHARED_CB_DATA Common;")
                                          .Replace("SIGMA_SHARED_CB_DATA", "public SIGMA_SHARED_CB_DATA Common;")
                                          .Replace(",", "")
                                          .Replace("NRD_CONSTANT(", "public")
                                          .Replace(") \\", ";")
                                          .Replace(")", ";");

                    structDef = RemoveWhitespaceLines(structDef);
                    writer.WriteLine(structDef);
                    writer.WriteLine();
                }

                {
                    List<string> resourcesDef = new List<string>();
                    resourcesDef.AddRange(RegexMatchesStartEndMultiline(text, "NRD_INPUT_TEXTURE_START", "NRD_INPUT_TEXTURE_END"));
                    resourcesDef.AddRange(RegexMatchesStartEndMultiline(text, "NRD_OUTPUT_TEXTURE_START", "NRD_OUTPUT_TEXTURE_END"));

                    HashSet<string> memberVars = new HashSet<string>();

                    if (resourcesDef.Count > 0)
                    {
                        string structName = file.Name.Replace(".resources.hlsli", "_ResourceSlot");
                        writer.WriteLine("static class " + structName +" {");

                        foreach (var res in resourcesDef)
                        {
                            foreach (var def in Regex.Matches(res, @"\(.*\)").Cast<Match>())
                            {
                                string item = def.Value.Substring(1, def.Value.Length - 2);
                                string[] desc = item.Split(',');
                                if (desc.Length == 4)
                                {
                                    string varName = desc[1].Trim();
                                    if (memberVars.Contains(varName))
                                        continue;

                                    writer.WriteLine(string.Format("    public static int {0} = UnityEngine.Shader.PropertyToID(\"{0}\");", varName));
                                    memberVars.Add(varName);
                                }
                            }
                        }
                        writer.WriteLine("}");
                    }
                }

                writer.WriteLine();
            }

            writer.WriteLine("}");
            writer.Flush();
            writer.Close();
        }

        static string[] RegexMatches(string input, string pattern)
        {
            return Regex.Matches(input, pattern, RegexOptions.Multiline).Select(m => m.Value).ToArray();
        }

        static string RegexMatch(string input, string pattern)
        {
            return Regex.Match(input, pattern, RegexOptions.Multiline).Value;
        }

        static string[] RegexMatchesStartEndMultiline(string input, string startWith, string endWith)
        {
            string pattern = string.Format(@"{0}((.|\n)*?){1}", startWith, endWith);
            return Regex.Matches(input, pattern, RegexOptions.Multiline).Select(m => m.Value).ToArray();
        }

        static string RemoveFirstLine(string s)
        {
            return s.Substring(s.IndexOf(Environment.NewLine, StringComparison.Ordinal) + 1);
        }

        static string RemoveWhitespaceLines(string s)
        {
            return Regex.Replace(s, @"^\s+$[\r\n]*", string.Empty, RegexOptions.Multiline);
        }

        const string nrdShaderFolder = "Shaders/NRDPass/";

        const string nrdCompilerConfiguration = @"
#ifndef NRD_CS_MAIN
    #define NRD_CS_MAIN main
#endif

#define NRD_EXPORT

#define NRD_CONSTANTS_START                                                         cbuffer globalConstants {
#define NRD_CONSTANT( constantType, constantName )                                  constantType constantName;
#define NRD_CONSTANTS_END                                                           };

#define NRD_INPUT_TEXTURE_START
#define NRD_INPUT_TEXTURE( resourceType, resourceName, regName, bindingIndex )      resourceType resourceName;
#define NRD_INPUT_TEXTURE_END

#define NRD_OUTPUT_TEXTURE_START
#define NRD_OUTPUT_TEXTURE( resourceType, resourceName, regName, bindingIndex )     resourceType resourceName;
#define NRD_OUTPUT_TEXTURE_END

#define NRD_SAMPLER_START
#define NRD_SAMPLER( resourceType, resourceName, regName, bindingIndex )            resourceType resourceName;
#define NRD_SAMPLER_END

#define gNearestClamp gPointClamp
#define gNearestMirror gPointMirror
";

        static string reblurSharedCBData
        {
            get
            {
                string text = File.ReadAllText(nrdShaderFolder + "Include/REBLUR/REBLUR_Config.hlsli");
                string desc = RegexMatch(text, @"#define REBLUR_SHARED_CB_DATA((.|\n)*?)^\s*$");
                desc = "public struct REBLUR_SHARED_CB_DATA {\n" + RemoveFirstLine(desc) + " };";
                desc = desc.Replace(",", "")
                           .Replace("NRD_CONSTANT(", "public")
                           .Replace(") \\", ";")
                           .Replace(")", ";");
                return desc;
            }
        }

        static string relaxSharedCBData
        {
            get
            {
                string text = File.ReadAllText(nrdShaderFolder + "Include/RELAX/RELAX_Config.hlsli");
                string desc = RegexMatch(text, @"#define RELAX_SHARED_CB_DATA((.|\n)*?)^\s*$");
                desc = "public struct RELAX_SHARED_CB_DATA {\n" + RemoveFirstLine(desc) + " };";
                desc = desc.Replace(",", "")
                           .Replace("NRD_CONSTANT(", "public")
                           .Replace(") \\", ";")
                           .Replace(")", ";");
                return desc;
            }
        }

        static string sigmaSharedCBData
        {
            get
            {
                string text = File.ReadAllText(nrdShaderFolder + "Include/SIGMA/SIGMA_Config.hlsli");
                string desc = RegexMatch(text, @"#define SIGMA_SHARED_CB_DATA((.|\n)*?)^\s*$");
                desc = "public struct SIGMA_SHARED_CB_DATA {\n" + RemoveFirstLine(desc) + " };";
                desc = desc.Replace(",", "")
                           .Replace("NRD_CONSTANT(", "public")
                           .Replace(") \\", ";")
                           .Replace(")", ";");
                return desc;
            }
        }
    }
}