Crafter.Graphics/shaders/ui-text.comp.glsl

//SPDX-License-Identifier: LGPL-3.0-only
//SPDX-FileCopyrightText: Copyright (C) 2026 Catcrafts®

#version 460
#extension GL_GOOGLE_include_directive : enable
#extension GL_KHR_shader_subgroup_basic : enable
#extension GL_KHR_shader_subgroup_arithmetic : enable
#include "ui-shared.glsl"

// One workgroup per 8×8 screen tile. The workgroup cooperatively streams the
// GlyphItem list in chunks of 64 (see ui-shared.glsl), culling each chunk
// against the tile and compacting survivors — in buffer order — into shared
// memory; every thread then accumulates over only those survivors.
layout(push_constant) uniform PC {
    UIDispatchHeader hdr;
    uint  fontTextureSlot;
    uint  fontSamplerSlot;
    uint  _p0;
    uint  _p1;
} pc;

layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;

shared vec4 s_rect[UI_CHUNK];
shared vec4 s_uv[UI_CHUNK];
shared vec4 s_color[UI_CHUNK];
// Per-glyph constants precomputed once at cooperative-load time and read by the
// per-pixel inner loop, instead of being recomputed for every pixel × glyph:
//   s_invRectSize = 1.0/rect.zw  (turns the per-pixel divide into a multiply)
//   s_ramp        = the SDF edge ramp (ui-shared.glsl's uiGlyphRamp), which
//                   depends only on the glyph's uv span and rect size
shared vec2 s_invRectSize[UI_CHUNK];
shared UIGlyphRamp s_ramp[UI_CHUNK];
shared uint s_keep[UI_CHUNK];
shared uint s_order[UI_CHUNK];
shared uint s_count;
shared uint s_subTotals[UI_CHUNK];  // per-subgroup survivor totals (carry step)

// Stable in-order compaction of one chunk's survivors via a per-subgroup
// exclusive prefix-sum plus a carry across subgroups, so a survivor's slot in
// s_order[] equals the number of survivors with a smaller local index — buffer
// (draw) order preserved. Replaces the old lane-0 serial scan. See
// ui-fused.comp.glsl for the full rationale; the body is identical.
void uiCompactChunk() {
    uint lid       = gl_LocalInvocationIndex;
    uint keep      = s_keep[lid];
    uint subPrefix = subgroupExclusiveAdd(keep);
    uint subTotal  = subgroupAdd(keep);
    if (subgroupElect())
        s_subTotals[gl_SubgroupID] = subTotal;
    barrier();

    uint base  = 0u;
    uint total = 0u;
    for (uint s = 0u; s < gl_NumSubgroups; ++s) {
        uint t = s_subTotals[s];
        if (s < gl_SubgroupID) base += t;
        total += t;
    }
    if (lid == 0u) s_count = total;

    if (keep != 0u)
        s_order[base + subPrefix] = lid;
}

void main() {
    ivec2 screenPx;
    bool  valid = uiResolveScreenPixel(pc.hdr, screenPx);

    // Defer the destination read-modify-write: a sparse UI leaves most tiles
    // untouched, so only load the pixel when the first surviving glyph blends
    // over it (`loaded`), and only store when something actually touched it.
    // The fused kernel amortizes a single load/store across all categories; the
    // standalone Dispatch* path has no such umbrella and otherwise pays a full
    // read-modify-write per empty tile.
    vec4 dst    = vec4(0.0);
    vec2 sp     = vec2(0.0);
    bool loaded = false;
    if (valid) sp = vec2(screenPx) + 0.5;

    vec2 tileMin, tileMax;
    uiTileBounds(tileMin, tileMax);

    // DispatchText draws only with a FontAtlas registered, so the font slots
    // always name it here.
    vec2 atlasSize = uiGlyphAtlasSize(pc.fontTextureSlot, pc.fontSamplerSlot);

    uint lid = gl_LocalInvocationIndex;
    for (uint base = 0u; base < pc.hdr.itemCount; base += UI_CHUNK) {
        uint idx  = base + lid;
        bool keep = false;
        if (idx < pc.hdr.itemCount) {
            s_rect[lid]  = uiGlyphHeap[pc.hdr.itemBuffer].items[idx].rect;
            s_uv[lid]    = uiGlyphHeap[pc.hdr.itemBuffer].items[idx].uv;
            s_color[lid] = uiGlyphHeap[pc.hdr.itemBuffer].items[idx].color;
            s_invRectSize[lid] = 1.0 / s_rect[lid].zw;
            s_ramp[lid]  = uiGlyphRamp(s_uv[lid], s_invRectSize[lid], atlasSize);
            keep = uiAabbOverlapsTile(s_rect[lid].xy, s_rect[lid].xy + s_rect[lid].zw,
                                      tileMin, tileMax);
        }
        s_keep[lid] = keep ? 1u : 0u;
        barrier();

        uiCompactChunk();
        barrier();

        if (valid) {
            for (uint j = 0u; j < s_count; ++j) {
                uint c = s_order[j];

                vec2 lo = s_rect[c].xy;
                vec2 hi = s_rect[c].xy + s_rect[c].zw;
                if (sp.x < lo.x || sp.y < lo.y) continue;
                if (sp.x >= hi.x || sp.y >= hi.y) continue;

                vec2 t  = (sp - s_rect[c].xy) * s_invRectSize[c];
                vec2 uv = mix(s_uv[c].xy, s_uv[c].zw, t);

                // Font slots are push constants, dynamically uniform as
                // uiGlyphCoverage needs.
                float a = uiGlyphCoverage(pc.fontTextureSlot, pc.fontSamplerSlot, uv, s_ramp[c], atlasSize);
                if (a <= 0.0) continue;

                vec4 col = s_color[c];
                vec4 src = vec4(col.rgb, col.a * a);
                if (!loaded) {
                    dst    = imageLoad(uiImages[pc.hdr.outImage], screenPx);
                    loaded = true;
                }
                dst = uiBlendOver(dst, src);
            }
        }
        barrier();
    }

    if (loaded) imageStore(uiImages[pc.hdr.outImage], screenPx, dst); // loaded ⇒ valid
}