GhostEngine/Ghost.DSL/ShaderCompiler/DSLShaderCompiler.cs

using Ghost.Core;
using Ghost.Core.Graphics;
using Ghost.DSL.ShaderParser;
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 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(ReadOnlySpan<PropertyDescriptor> properties)
    {
        if (properties.IsEmpty)
        {
            return 0;
        }

        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,
            hlsl = semantics.hlsl
        };

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

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

        descriptor.globalProperties = shaderGlobalProperties ?? Array.Empty<PropertyDescriptor>();
        descriptor.properties = shaderLocalProperties ?? Array.Empty<PropertyDescriptor>();
        descriptor.cbufferSize = CalculateCBufferSize(descriptor.properties);

        if (semantics.passes != null)
        {
            descriptor.passes = new PassDescriptor[semantics.passes.Count];
            for (int i = 0; i < semantics.passes.Count; i++)
            {
                var pass = semantics.passes[i];
                var localPipeline = MeragePipeline(pass.localPipeline, PipelineState.Default);
                descriptor.passes[i] = new PassDescriptor
                {
                    identifier = GetPassUniqueId(semantics, pass),
                    name = pass.name,
                    taskShader = pass.taskShader,
                    meshShader = pass.meshShader,
                    pixelShader = pass.pixelShader,
                    localPipeline = localPipeline,
                    defines = pass.defines?.ToArray() ?? Array.Empty<string>(),
                    includes = pass.includes?.ToArray() ?? Array.Empty<string>(),
                    keywords = pass.keywords?.ToArray() ?? Array.Empty<KeywordsGroup>(),
                    hlsl = pass.hlsl
                };
            }
        }
        else
        {
            descriptor.passes = Array.Empty<PassDescriptor>();
        }

        return descriptor;
    }

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

            // Use ANTLR4 parser
            var shaderModels = AntlrShaderCompiler.ParseShaders(source, out var parseErrors);

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

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

            if (shaderModels.Count == 0)
            {
                return Result.Failure("No shader found in the provided file.");
            }

            // Convert to semantics
            var model = AntlrShaderCompiler.ConvertToSemantics(shaderModels[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 (ref var pass in desc.passes.AsSpan())
            {
                if (pass.includes == null)
                {
                    pass.includes = new string[2];
                }
                else
                {
                    Array.Resize(ref pass.includes, pass.includes.Length + 2);
                    // Shift existing includes to make room for the two new includes at the front.
                    pass.includes.AsSpan(0, pass.includes.Length - 2).CopyTo(pass.includes.AsSpan(2));
                }

                pass.includes[0] = globalPropResult.Value;
                pass.includes[1] = 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.Graphics/Shaders/Includes/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(ReadOnlySpan<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.Graphics/Shaders/Includes/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;
    }
}