#version 450

// Output is a global indices[] buffer, indexed by gl_InstanceIndex.

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

layout(std430, set = 0, binding = 0) readonly buffer SortedBlocks
{
    uint blocks[];
} sortedBlocks;

layout(std430, set = 0, binding = 1) writeonly buffer DrawIndices
{
    uint indices[];
} outIndices;

layout(push_constant) uniform Push
{
    uvec4 header; // x=base, y=count, z=flags (bit0=identity)
} pc;

const uint BLOCK_SIZE = 256u;

void main()
{
    uint i = gl_GlobalInvocationID.x;
    uint count = pc.header.y;
    if (i >= count) return;

    uint base = pc.header.x;
    uint flags = pc.header.z;

    uint outIdx = base + i;
    uint particleIdx = base + i;

    if ((flags & 1u) == 0u)
    {
        uint blockRank = i / BLOCK_SIZE;
        uint within = i - blockRank * BLOCK_SIZE;
        uint block = sortedBlocks.blocks[blockRank];
        particleIdx = base + block * BLOCK_SIZE + within;
    }

    outIndices.indices[outIdx] = particleIdx;
}