jump_flood.comp

#version 460

const uint UINT_MAX = 4294967295U;
const ivec2 normalizedOffsets[9] = {
    { -1, -1 }, { 0, -1 }, { 1, -1 },
    { -1,  0 }, { 0,  0 }, { 1,  0 },
    { -1,  1 }, { 0,  1 }, { 1,  1 }
};

layout (set = 0, binding = 0, rg16ui) uniform uimage2DArray pingPongImage;

layout (push_constant, std430) uniform PushConstant {
    bool forward;
    uint sampleOffset;
} pc;

layout (local_size_x = 16, local_size_y = 16) in;

uint length2(uvec2 v) {
    return v.x * v.x + v.y * v.y;
}

void main(){
    uvec2 imageExtent = imageSize(pingPongImage).xy;
    if (gl_GlobalInvocationID.x >= imageExtent.x || gl_GlobalInvocationID.y >= imageExtent.y) {
        return;
    }

    uvec2 closestSeedCoord;
    uint closestSeedDistanceSq = UINT_MAX;
    for (uint i = 0; i < 9; ++i){
        uvec2 seedCoord = imageLoad(pingPongImage, ivec3(ivec2(gl_GlobalInvocationID.xy) + int(pc.sampleOffset) * normalizedOffsets[i], !pc.forward)).xy;
        if (seedCoord != uvec2(0U)) {
            uint seedDistanceSq = length2(seedCoord - gl_GlobalInvocationID.xy);
            if (seedDistanceSq < closestSeedDistanceSq) {
                closestSeedDistanceSq = seedDistanceSq;
                closestSeedCoord = seedCoord;
            }
        }
    }

    imageStore(pingPongImage, ivec3(ivec2(gl_GlobalInvocationID.xy), pc.forward), uvec4(closestSeedDistanceSq == UINT_MAX ? uvec2(0) : closestSeedCoord, 0, 0));
}