GhostEngine/Ghost.Shader/Compiler/SDLCompiler.cs

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

namespace Ghost.SDL.Compiler;

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

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

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

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

        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 SDLSemantics? SemanticAnalysis(SDLSyntax syntax, out List<SDLError> errors)
    {
        errors = new();

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

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

        var sortedList = new List<SDLSemantics>();
        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(SDLSemantics shader, PassSemantic pass)
    {
        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(SDLSemantics 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 (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;
    }

    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.Fail("Failed to compile shader due to errors:\n" + errorMessages.ToString());
            }

            var desc = ResolveShader(model);
            var globalPropPath = GenerateGlobalProperties(desc.globalProperties, generatedOutputDirectory);

            foreach (var pass in desc.passes)
            {
                if (pass is not FullPassDescriptor fullPass)
                {
                    continue;
                }

                fullPass.includes ??= new List<string>();
                fullPass.includes.Add(globalPropPath);
                fullPass.generatedCodePath = GeneratePass(fullPass, generatedOutputDirectory);
            }

            return desc;
        }
        catch (Exception ex)
        {
            return Result.Fail("Failed to generate shader files: " + 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.Texture2DBindless => "TEXTURE2D_BINDLESS",
            ShaderPropertyType.Texture3DBindless => "TEXTURE3D_BINDLESS",
            ShaderPropertyType.TextureCubeBindless => "TEXTURECUBE_BINDLESS",
            ShaderPropertyType.Texture2DArrayBindless => "TEXTURE2D_ARRAY_BINDLESS",
            ShaderPropertyType.TextureCubeArrayBindless => "TEXTURECUBE_ARRAY_BINDLESS",
            _ => throw new ArgumentOutOfRangeException(nameof(type), $"Unsupported shader property type: {type}")
        };
    }

    public static string GeneratePass(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.includes != null)
        {
            foreach (var include in fullPass.includes)
            {
                sb.Append($@"
#include ""{include}""");
            }

            sb.AppendLine();
        }

        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 string GenerateGlobalProperties(List<PropertyDescriptor> globalProperties, string targetDirectory)
    {
        if (!Directory.Exists(targetDirectory))
        {
            throw new ArgumentException("Target directory does not exist.", nameof(targetDirectory));
        }

        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 globalProperties)
        {
            sb.Append($@"
{ShaderPropertyTypeToHLSLType(prop.type)} {prop.name};");
        }
        sb.AppendLine(@"
};

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

        return globalFilePath;
    }
}