GhostEngine/Ghost.Shader/Compiler/ShaderCompiler.cs

using Ghost.Core.Graphics;
using Ghost.Shader.Compiler.Parser;
using System.Text;

namespace Ghost.Shader.Compiler;

public struct ShaderError
{
    public string message;
    public int line;
    public int column;

    public readonly override string ToString()
    {
        return $"Error at {line}:{column} - {message}";
    }
}

internal static class ShaderCompiler
{
    private const string _GLOBAL_PROPERTY_FILE_NAME = "GlobalData.g.hlsl";
    private const string _GENERATED_FILE_HEADER = "// Auto-generated shader file. Please do not edit this file directly.";

    private struct ShaderInheritance
    {
        public ShaderSemantics? parent;
        public List<ShaderInheritance>? children;
    }

    public static List<ShaderSyntax> ParseShaders(TokenStream stream)
    {
        var shaders = new List<ShaderSyntax>();

        while (stream.TryPeek(out var nextToken))
        {
            if (ShaderBlock.ShouldEnter(nextToken))
            {
                var shader = ShaderBlock.Parse(stream.SliceNextBlock());
                shaders.Add(shader);
            }
            else if (nextToken.Match(TokenType.EndOfFile))
            {
                stream.Consume();
            }
            else
            {
                throw new Exception($"Unexpected token '{nextToken.lexeme}' at top level. Expected 'shader' declaration.");
            }
        }

        return shaders;
    }

    public static ShaderSemantics? SemanticAnalysis(ShaderSyntax syntax, out List<ShaderError> errors)
    {
        errors = new();

        if (string.IsNullOrWhiteSpace(syntax.name.lexeme))
        {
            errors.Add(new ShaderError
            {
                message = "Shader name cannot be empty.",
                line = syntax.name.line,
                column = syntax.name.column
            });

            return null;
        }

        var shaderModel = ShaderBlock.SemanticAnalysis(syntax, errors);
        return shaderModel;
    }

    private static List<ShaderSemantics>? TopologicalSort(ReadOnlySpan<ShaderSemantics> semantics)
    {
        var inDegrees = new Dictionary<string, int>();
        var childrenMap = new Dictionary<string, List<string>>();
        var semanticsMap = new Dictionary<string, ShaderSemantics>();

        foreach (var s in semantics)
        {
            inDegrees[s.name] = 0;
            childrenMap[s.name] = new List<string>();
            semanticsMap[s.name] = s;
        }

        foreach (var s in semantics)
        {
            if (!string.IsNullOrEmpty(s.fallback) && semanticsMap.ContainsKey(s.fallback))
            {
                childrenMap[s.fallback].Add(s.name);
                inDegrees[s.name]++;
            }
        }

        var queue = new Queue<ShaderSemantics>();
        foreach (var s in semantics)
        {
            if (inDegrees[s.name] == 0)
            {
                queue.Enqueue(s);
            }
        }

        var sortedList = new List<ShaderSemantics>();
        while (queue.Count > 0)
        {
            var current = queue.Dequeue();
            sortedList.Add(current);

            foreach (var childName in childrenMap[current.name])
            {
                inDegrees[childName]--;
                if (inDegrees[childName] == 0)
                {
                    queue.Enqueue(semanticsMap[childName]);
                }
            }
        }

        // If there's a cycle, the graph will not be fully traversed.
        return sortedList.Count == semantics.Length ? sortedList : null;
    }

    private static string GetPassUniqueId(ShaderSemantics shader, PassSemantic pass)
    {
        //static ulong Fnv1a64(ReadOnlySpan<char> data)
        //{
        //    const ulong offset = 14695981039346656037;
        //    const ulong prime = 1099511628211;

        //    var hash = offset;

        //    foreach (var b in data)
        //    {
        //        hash ^= b;
        //        hash *= prime;
        //    }

        //    return hash;
        //}

        //return $"{Fnv1a64(shader.name)}_{pass.name}";

        return $"{shader.name}_{pass.name}";
    }

    private static PipelineDescriptor MeragePipeline(PipelineSemantic? semantic, PipelineDescriptor parent)
    {
        if (semantic == null)
        {
            return parent;
        }

        return new PipelineDescriptor
        {
            zTest = semantic.zTest ?? parent.zTest,
            zWrite = semantic.zWrite ?? parent.zWrite,
            cull = semantic.cull ?? parent.cull,
            blend = semantic.blend ?? parent.blend,
            colorMask = semantic.colorMask ?? parent.colorMask
        };
    }

    private static List<PropertyDescriptor> MergeProperties(List<PropertySemantic>? semantics, List<PropertyDescriptor>? parent)
    {
        var result = new List<PropertyDescriptor>();
        if (parent != null)
        {
            result.AddRange(parent);
        }

        if (semantics != null)
        {
            foreach (var prop in semantics)
            {
                if (prop.scope == PropertyScope.Local)
                {
                    result.Add(new PropertyDescriptor
                    {
                        name = prop.name,
                        type = prop.type,
                        defaultValue = prop.defaultValue
                    });
                }
            }
        }

        return result.DistinctBy(p => p.name).ToList();
    }

    // TODO: Implement shader inheritance resolution, including property and pass merging.
    // Currently, we just ignore inheritance.
    public static ShaderDescriptor ResolveShader(ShaderSemantics semantics)
    {
        var descriptor = new ShaderDescriptor
        {
            name = semantics.name
        };

        var shaderGlobalProperties = semantics.properties?.Where(p => p.scope == PropertyScope.Global).Select(p => new PropertyDescriptor
        {
            name = p.name,
            type = p.type,
            defaultValue = p.defaultValue
        }).ToList();

        var shaderLocalProperties = semantics.properties?.Where(p => p.scope == PropertyScope.Local).Select(p => new PropertyDescriptor
        {
            name = p.name,
            type = p.type,
            defaultValue = p.defaultValue
        }).ToList();

        if (shaderGlobalProperties != null)
        {
            descriptor.globalProperties.AddRange(shaderGlobalProperties);
        }

        if (semantics.passes != null)
        {
            foreach (var pass in semantics.passes)
            {
                var localPipeline = MeragePipeline(pass.localPipeline, PipelineDescriptor.Default);
                var localProperties = MergeProperties(pass.localProperties, shaderLocalProperties); // TODO: Merge with base shader properties if inheritance is implemented.

                var fullPass = new FullPassDescriptor
                {
                    uniqueIdentifier = GetPassUniqueId(semantics, pass),
                    taskShader = pass.taskShader,
                    meshShader = pass.meshShader,
                    pixelShader = pass.pixelShader,
                    localPipeline = localPipeline,
                    defines = pass.defines,
                    includes = pass.includes,
                    keywords = pass.keywords,
                    properties = localProperties
                };

                descriptor.passes.Add(fullPass);
            }
        }

        return descriptor;
    }

    private static string ShaderPropertyTypeToHLSLType(ShaderPropertyType type)
    {
        return type switch
        {
            ShaderPropertyType.Float => "float",
            ShaderPropertyType.Float2 => "float2",
            ShaderPropertyType.Float3 => "float3",
            ShaderPropertyType.Float4 => "float4",
            ShaderPropertyType.Int => "int",
            ShaderPropertyType.Int2 => "int2",
            ShaderPropertyType.Int3 => "int3",
            ShaderPropertyType.Int4 => "int4",
            ShaderPropertyType.UInt => "uint",
            ShaderPropertyType.UInt2 => "uint2",
            ShaderPropertyType.UInt3 => "uint3",
            ShaderPropertyType.UInt4 => "uint4",
            ShaderPropertyType.Bool => "bool",
            ShaderPropertyType.Bool2 => "bool2",
            ShaderPropertyType.Bool3 => "bool3",
            ShaderPropertyType.Bool4 => "bool4",
            // NOTE: Textures here are bindless, represented as uint (descriptor index).
            ShaderPropertyType.Texture2D => "TEXTURE2D_BINDLESS",
            ShaderPropertyType.Texture3D => "TEXTURE3D_BINDLESS",
            ShaderPropertyType.TextureCube => "TEXTURECUBE_BINDLESS",
            ShaderPropertyType.Texture2DArray => "TEXTURE2D_ARRAY_BINDLESS",
            ShaderPropertyType.TextureCubeArray => "TEXTURECUBE_ARRAY_BINDLESS",
            _ => throw new ArgumentOutOfRangeException(nameof(type), $"Unsupported shader property type: {type}")
        };
    }

    public static string CompilePass(IPassDescriptor descriptor, string targetDirectory)
    {
        if (descriptor is not FullPassDescriptor fullPass)
        {
            throw new NotSupportedException("Only full pass descriptors are supported for compilation.");
        }

        if (!Directory.Exists(targetDirectory))
        {
            throw new ArgumentException("Target directory does not exist.", nameof(targetDirectory));
        }

        var outputFileName = fullPass.uniqueIdentifier.Replace(' ', '_');
        var outputFilePath = Path.Combine(targetDirectory, outputFileName + ".g.hlsl");
        var outputDirectory = Path.GetDirectoryName(outputFilePath);

        if (!Directory.Exists(outputDirectory))
        {
            Directory.CreateDirectory(outputDirectory!);
        }

        using var fileStream = File.CreateText(outputFilePath);
        var fileDefine = outputFileName.Replace('/', '_').ToUpperInvariant() + "_G_HLSL";

        var sb = new StringBuilder();

        sb.AppendLine(_GENERATED_FILE_HEADER);
        sb.AppendLine(@$"
#ifndef {fileDefine}
#define {fileDefine}

#include ""F:/csharp/GhostEngine/Ghost.Shader/BuiltIn/Common.hlsl""");

        if (fullPass.properties != null)
        {
            sb.Append(@"
struct PerMaterialData
{");
            foreach (var prop in fullPass.properties)
            {
                sb.Append($@"
    {ShaderPropertyTypeToHLSLType(prop.type)} {prop.name};");
            }
            sb.Append(@"
};");
        }

        sb.AppendLine();
        sb.AppendLine(@$"
#endif // {fileDefine}");

        fileStream.Write(sb.ToString());

        return outputFilePath;
    }

    public static void CompileShader(ShaderDescriptor descriptor, string targetDirectory)
    {
        if (!Directory.Exists(targetDirectory))
        {
            throw new ArgumentException("Target directory does not exist.", nameof(targetDirectory));
        }

        // Generate global property file.
        if (descriptor.globalProperties.Count > 0)
        {
            var globalFilePath = Path.Combine(targetDirectory, _GLOBAL_PROPERTY_FILE_NAME);
            using var globalFileStream = File.CreateText(globalFilePath);

            var sb = new StringBuilder();

            sb.AppendLine(_GENERATED_FILE_HEADER);
            sb.Append(@"
#ifndef GLOBALDATA_G_HLSL
#define GLOBALDATA_G_HLSL

#include ""F:/csharp/GhostEngine/Ghost.Shader/BuiltIn/Common.hlsl""

struct GlobalData
{");
            foreach (var prop in descriptor.globalProperties)
            {
                sb.Append($@"
    {ShaderPropertyTypeToHLSLType(prop.type)} {prop.name};");
            }
            sb.AppendLine(@"
};

#endif // GLOBALDATA_G_HLSL");
            globalFileStream.Write(sb.ToString());
        }

        // Compile each pass.
        foreach (var pass in descriptor.passes)
        {
            CompilePass(pass, targetDirectory);
        }
    }
}