acdream/tools/ShaderCompiler/VulkanGlslPreamble.cs

using System.Text;

namespace AcDream.Tools.ShaderCompiler;

/// <summary>
/// Campaign V slice V6c, plan §3.4 and §4.6: the Vulkan half of acdream's
/// dual-dialect GLSL.
///
/// <para>The GLSL sources under <c>Rendering/Shaders</c> are the single source of
/// truth for both backends, and they are written in the dialect GL accepts. The
/// three things Vulkan needs on top of that are not source edits — they are
/// definitions the compiler injects immediately after the <c>#version</c> line:
/// </para>
///
/// <list type="number">
/// <item><c>ACDREAM_UBO_SET</c> becomes <c>set = 1,</c>. Under GL it expands to
/// nothing, because GL keeps the SSBO and UBO binding namespaces separate and
/// <c>BatchBuffer</c> (SSBO binding 1) and <c>SceneLighting</c> (UBO binding 1)
/// can share a number. Vulkan has one namespace per set, so moving uniform
/// buffers to set 1 preserves both numbers. <c>common.glsl</c> has carried this
/// macro since V2 for exactly this moment.</item>
/// <item>The texture table becomes a real descriptor array at set 2 binding 0,
/// and <c>ACDREAM_TEXTURE_HANDLE</c> becomes an index rather than a packed
/// bindless handle. <c>nonuniformEXT</c> is required, not optional: within one
/// multi-draw dispatch different draws read different <c>Batches[]</c> entries,
/// and "dynamically uniform" is defined over the whole dispatch on some
/// implementations.</item>
/// <item>The shared 96-byte push-constant block is declared, and each loose
/// uniform name is <c>#define</c>d onto its member. Vulkan GLSL has no default
/// uniform block, so this is the only way the same source can declare
/// <c>uniform mat4 uViewProjection;</c> for GL and read a push constant for
/// Vulkan.</item>
/// </list>
///
/// <para>Injection rather than source rewriting is the whole design. A textual
/// transform over the real shader bodies would be a small compiler with its own
/// failure modes, and it would fail silently — a shader that compiles but reads
/// the wrong storage buffer looks exactly like a shader that works until
/// someone renders with it. A preamble either defines what the body needs or the
/// compile fails loudly, which is a property worth more than convenience.</para>
/// </summary>
internal static class VulkanGlslPreamble
{
    /// <summary>
    /// The loose uniform names the shared push-constant block carries, in the
    /// order <c>GpuPushConstants</c> declares them. Any shader whose uniforms are
    /// all in this list needs nothing but the preamble; any shader with a
    /// uniform outside it cannot be expressed against the pinned contract and is
    /// reported rather than guessed at.
    /// </summary>
    internal static IReadOnlyList<string> PushConstantFields { get; } =
    [
        "uViewProjection",
        "uDrawIDOffset",
        "uLightingMode",
        "uRenderPass",
        "uLightDebug",
        "uTextureIndexA",
        "uTextureIndexB",
        "uParamA",
        "uParamB",
    ];

    /// <summary>
    /// Builds the text inserted after the <c>#version</c> directive.
    /// <paramref name="stage"/> only affects which stage-specific rewrites are
    /// emitted.
    /// </summary>
    internal static string Build(string stage)
    {
        var text = new StringBuilder();
        text.AppendLine("// ---- injected by tools/compile-shaders.ps1 (Campaign V slice V6c) ----");
        text.AppendLine("// Vulkan dialect only. The GL backend compiles the same source with none");
        text.AppendLine("// of this, which is what keeps one GLSL file the single source of truth.");
        text.AppendLine("#extension GL_EXT_nonuniform_qualifier : require");
        if (stage == "vert")
        {
            // gl_DrawID is Vulkan's shaderDrawParameters feature, and glslang
            // still gates the identifier behind the ARB extension name even when
            // targeting Vulkan. Declared here so a source that does not name it
            // still gets it.
            text.AppendLine("#extension GL_ARB_shader_draw_parameters : require");
        }

        text.AppendLine();
        text.AppendLine("// §3.4 set 1: every uniform buffer. Under GL this macro expands to nothing.");
        text.AppendLine("#undef ACDREAM_UBO_SET");
        text.AppendLine("#define ACDREAM_UBO_SET set = 1,");
        text.AppendLine();
        text.AppendLine("// §4.4 set 2: the global sampled-texture table that replaces");
        text.AppendLine("// GL_ARB_bindless_texture. Variable count, partially bound,");
        text.AppendLine("// update-after-bind; the CPU never writes it per frame.");
        text.AppendLine("layout(set = 2, binding = 0) uniform sampler2DArray uTextures[];");
        text.AppendLine("#undef ACDREAM_TEXTURE_HANDLE");
        text.AppendLine("#define ACDREAM_TEXTURE_HANDLE(idx) (idx)");
        text.AppendLine("#define ACDREAM_TEXTURE(idx) uTextures[nonuniformEXT(uint(idx))]");
        // Slice V6d: the 2-D read. GL reconstructs a sampler2D from the entry's
        // bindless handle; Vulkan has one descriptor array whose element type is
        // fixed at sampler2DArray, so a 2-D entry is a one-layer array read at
        // layer 0. See common.glsl for the GL half and for why the UI's
        // textures stay plain GL_TEXTURE_2D objects.
        text.AppendLine(
            "#define ACDREAM_SAMPLE_2D(idx, uv) texture(ACDREAM_TEXTURE(idx), vec3((uv), 0.0))");
        text.AppendLine();
        text.AppendLine("// §3.4 push constants: one shared 96-byte block, so switching pipelines");
        text.AppendLine("// mid-pass invalidates neither descriptors nor constants.");
        text.AppendLine("layout(push_constant) uniform AcdreamPushBlock {");
        text.AppendLine("    mat4 viewProjection;");
        text.AppendLine("    int drawIdOffset;");
        text.AppendLine("    int lightingMode;");
        text.AppendLine("    int renderPass;");
        text.AppendLine("    int lightDebug;");
        text.AppendLine("    uint textureIndexA;");
        text.AppendLine("    uint textureIndexB;");
        text.AppendLine("    float paramA;");
        text.AppendLine("    float paramB;");
        text.AppendLine("} acdreamPush;");
        text.AppendLine();
        text.AppendLine("#define uViewProjection acdreamPush.viewProjection");
        text.AppendLine("#define uDrawIDOffset   acdreamPush.drawIdOffset");
        text.AppendLine("#define uLightingMode   acdreamPush.lightingMode");
        text.AppendLine("#define uRenderPass     acdreamPush.renderPass");
        text.AppendLine("#define uLightDebug     acdreamPush.lightDebug");
        text.AppendLine("#define uTextureIndexA  acdreamPush.textureIndexA");
        text.AppendLine("#define uTextureIndexB  acdreamPush.textureIndexB");
        text.AppendLine("#define uParamA         acdreamPush.paramA");
        text.AppendLine("#define uParamB         acdreamPush.paramB");
        text.AppendLine();
        text.AppendLine("// §4.6: gl_DrawIDARB stays as written — glslang exposes it for Vulkan");
        text.AppendLine("// under the same ARB extension name. gl_InstanceIndex already includes");
        text.AppendLine("// firstInstance, so the GL idiom gl_BaseInstanceARB + gl_InstanceID");
        text.AppendLine("// collapses to it exactly.");
        text.AppendLine(stage == "vert"
            ? "#define gl_BaseInstanceARB 0\n"
                + "#define gl_InstanceID gl_InstanceIndex\n"
                + "#define gl_VertexID gl_VertexIndex"
            : "// (the vertex/instance-index rewrites apply to the vertex stage only)");
        text.AppendLine("// ---- end injected preamble ----");
        return text.ToString();
    }

    /// <summary>
    /// Returns <paramref name="source"/> with the preamble inserted after its
    /// <c>#version</c> line and the version raised to 450, which is the floor for
    /// Vulkan GLSL. Everything else is untouched: this never edits a shader body.
    /// </summary>
    internal static string Apply(string source, string stage)
    {
        ArgumentNullException.ThrowIfNull(source);
        string[] lines = source.Replace("\r\n", "\n").Split('\n');
        var output = new StringBuilder();
        bool injected = false;

        foreach (string line in lines)
        {
            string trimmed = line.TrimStart();
            if (!injected && trimmed.StartsWith("#version", StringComparison.Ordinal))
            {
                // 450 is the floor for Vulkan GLSL; a source already asking for
                // more keeps it, because a shader that opted into 460 did so for
                // a feature and quietly downgrading it would be a silent change.
                output.AppendLine(HighestVersion(trimmed) >= 460 ? "#version 460 core" : "#version 450 core");
                output.Append(Build(stage));
                injected = true;
                continue;
            }

            // Bindless textures are what the set-2 descriptor array replaces;
            // requiring the extension under Vulkan is an error rather than a
            // no-op. (GL_ARB_shader_draw_parameters is kept — glslang still gates
            // gl_DrawID behind that name when targeting Vulkan.)
            if (trimmed.StartsWith("#extension GL_ARB_bindless_texture", StringComparison.Ordinal))
            {
                output.AppendLine($"// (dropped for Vulkan: {trimmed})");
                continue;
            }

            // Vulkan GLSL has no default uniform block, so a loose
            // `uniform mat4 uViewProjection;` is illegal however it is spelled.
            // Dropping the DECLARATION is what lets the preamble's #define
            // redirect the name onto a push-constant member; a #define alone
            // would only rewrite the declaration into a worse one. Uniform BLOCK
            // declarations (which carry a `{`) are untouched, and an opaque
            // sampler or a name with no push-constant home simply becomes an
            // undeclared identifier — a loud, specific compiler error naming the
            // shader that still needs its port slice.
            if (IsDefaultBlockUniformDeclaration(trimmed))
            {
                output.AppendLine($"// (declaration dropped for Vulkan: {trimmed})");
                continue;
            }

            output.AppendLine(line);
        }

        if (!injected)
        {
            throw new InvalidOperationException(
                "The shader has no #version directive, so there is nowhere to inject the Vulkan preamble.");
        }

        return output.ToString();
    }

    /// <summary>The numeric version a <c>#version</c> directive asks for, or 450 when it cannot be read.</summary>
    internal static int HighestVersion(string versionDirective)
    {
        string[] parts = versionDirective.Split(
            [' ', '\t'],
            StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries);
        return parts.Length >= 2 && int.TryParse(parts[1], out int version) ? version : 450;
    }

    /// <summary>
    /// True for a loose (default-block) uniform declaration —
    /// <c>uniform mat4 uViewProjection;</c> or <c>uniform float uTexTiling[36];</c> —
    /// and false for a uniform BLOCK, which opens a brace and is legal in both
    /// dialects.
    /// </summary>
    internal static bool IsDefaultBlockUniformDeclaration(string trimmedLine)
    {
        ArgumentNullException.ThrowIfNull(trimmedLine);
        if (!trimmedLine.StartsWith("uniform ", StringComparison.Ordinal))
            return false;

        // A block declaration is `uniform Name {` — possibly with the brace on
        // the next line, in which case there is no semicolon here either.
        int comment = trimmedLine.IndexOf("//", StringComparison.Ordinal);
        string code = comment >= 0 ? trimmedLine[..comment] : trimmedLine;
        return !code.Contains('{') && code.TrimEnd().EndsWith(';');
    }
}