GhostEngine/Ghost.DSL/ShaderCompiler/DSLShaderCompiler.cs

using Ghost.Core;
using Ghost.Core.Graphics;
using Ghost.DSL.ShaderCompiler.Parser;
using System.Text;

namespace Ghost.DSL.ShaderCompiler;

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

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

internal static class DSLShaderCompiler
{
    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 DSLShaderSemantics? parent;
    //     public List<ShaderInheritance>? children;
    // }

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

        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 DSLShaderSemantics? SemanticAnalysis(DSLShaderSyntax syntax, out List<DSLShaderError> errors)
    {
        errors = new List<DSLShaderError>();

        if (string.IsNullOrWhiteSpace(syntax.name.lexeme))
        {
            errors.Add(new DSLShaderError
            {
                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<DSLShaderSemantics>? TopologicalSort(ReadOnlySpan<DSLShaderSemantics> semantics)
    {
        var inDegrees = new Dictionary<string, int>();
        var childrenMap = new Dictionary<string, List<string>>();
        var semanticsMap = new Dictionary<string, DSLShaderSemantics>();

        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<DSLShaderSemantics>();
        foreach (var s in semantics)
        {
            if (inDegrees[s.name] == 0)
            {
                queue.Enqueue(s);
            }
        }

        var sortedList = new List<DSLShaderSemantics>();
        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(DSLShaderSemantics shader, PassSemantic pass)
    {
        return $"{shader.name}_{pass.name}";
    }

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

        return new PipelineState
        {
            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 uint CalculateCBufferSize(List<PropertyDescriptor> properties)
    {
        var currentOffset = 0u;

        foreach (var prop in properties)
        {
            var size = prop.type.GetSize();

            if ((currentOffset % 16) + size > 16)
            {
                currentOffset = (currentOffset + 15u) & ~15u;
            }

            currentOffset += size;
        }

        return (currentOffset + 15u) & ~15u;
    }

    // TODO: Implement shader inheritance resolution, including property and pass merging.
    // Currently, we just ignore inheritance.
    public static ShaderDescriptor ResolveShader(DSLShaderSemantics 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 ??= new List<PropertyDescriptor>();
            descriptor.globalProperties.AddRange(shaderGlobalProperties);
        }

        if (shaderLocalProperties != null)
        {
            descriptor.properties ??= new List<PropertyDescriptor>();
            descriptor.properties.AddRange(shaderLocalProperties);
            descriptor.cbufferSize = CalculateCBufferSize(descriptor.properties);
        }

        if (semantics.passes != null)
        {
            foreach (var pass in semantics.passes)
            {
                var localPipeline = MeragePipeline(pass.localPipeline, PipelineState.Default);
                var fullPass = new PassDescriptor
                {
                    uniqueIdentifier = GetPassUniqueId(semantics, pass),
                    name = pass.name,
                    taskShader = pass.taskShader,
                    meshShader = pass.meshShader,
                    pixelShader = pass.pixelShader,
                    localPipeline = localPipeline,
                    defines = pass.defines,
                    keywords = pass.keywords,
                };

                descriptor.passes.Add(fullPass);
            }
        }

        return descriptor;
    }

    public static Result<ShaderDescriptor> CompileShader(string shaderPath, string generatedOutputDirectory)
    {
        try
        {
            var source = File.ReadAllText(shaderPath);

            var lexer = new Lexer(source);
            var stream = new TokenStream(lexer.Tokenize());
            var shaderInfo = ParseShaders(stream);
            var model = SemanticAnalysis(shaderInfo[0], out var errors);

            if (errors.Count != 0 || model == null)
            {
                var errorMessages = new StringBuilder();
                foreach (var error in errors)
                {
                    errorMessages.AppendLine(error.ToString());
                }

                return Result.Failure("Failed to compile shader due to errors:\n" + errorMessages.ToString());
            }

            var desc = ResolveShader(model);
            var globalPropResult = GenerateGlobalProperties(desc.globalProperties, generatedOutputDirectory);
            if (globalPropResult.IsFailure)
            {
                return Result.Failure("Failed to generate global properties: " + globalPropResult.Message);
            }

            var generatedResult = GenerateShaderCode(desc, generatedOutputDirectory);
            if (generatedResult.IsFailure)
            {
                return Result.Failure("Failed to generate pass files: " + generatedResult.Message);
            }

            foreach (var pass in desc.passes)
            {
                if (pass is PassDescriptor fullPass)
                {
                    fullPass.includes ??= new List<string>();
                    fullPass.includes.Add(globalPropResult.Value);
                    fullPass.includes.Add(generatedResult.Value);
                }
            }

            return desc;
        }
        catch (Exception ex)
        {
            return Result.Failure("Failed to compile shader: " + ex.Message);
        }
    }

    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",
            ShaderPropertyType.Texture3D => "TEXTURE3D",
            ShaderPropertyType.TextureCube => "TEXTURECUBE",
            ShaderPropertyType.Texture2DArray => "TEXTURE2D_ARRAY",
            ShaderPropertyType.TextureCubeArray => "TEXTURECUBE_ARRAY",
            ShaderPropertyType.Sampler => "SAMPLER",
            _ => throw new ArgumentOutOfRangeException(nameof(type), $"Unsupported shader property type: {type}")
        };
    }

    public static Result<string> GenerateShaderCode(ShaderDescriptor descriptor, string targetDirectory)
    {
        if (!Directory.Exists(targetDirectory))
        {
            return Result.Failure("Target directory does not exist.");
        }

        var outputFileName = descriptor.name.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.DSL/BuiltIn/Common.hlsl""");

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

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

        fileStream.Write(sb.ToString());

        return outputFilePath;
    }

    public static Result<string> GenerateGlobalProperties(List<PropertyDescriptor> globalProperties, string targetDirectory)
    {
        if (!Directory.Exists(targetDirectory))
        {
            return Result.Failure("Target directory does not exist.");
        }

        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.DSL/BuiltIn/Common.hlsl""

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

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

        return globalFilePath;
    }
}