#version 450
// ============================================================================
// Stream compaction scatter.
//
// Given:
// flags[N] - 1 if element is active, 0 otherwise
// scan_result[N] - exclusive prefix sum of flags[] (output index for active)
// source[N] - source data to compact (or use index as data)
//
// Produces:
// compacted[scan_result[i]] = source[i] for all i where flags[i] == 1
// count[0] = total number of active elements (= scan_result[N-1] + flags[N-1])
//
// Dispatch: ceil(N / 256), 1, 1
// ============================================================================
#define THREADS 256
layout(local_size_x = THREADS, local_size_y = 1, local_size_z = 1) in;
layout(set = 0, binding = 0, std430) readonly buffer FlagsBuf { uint flags[]; };
layout(set = 0, binding = 1, std430) readonly buffer ScanBuf { uint scan_result[]; };
layout(set = 0, binding = 2, std430) readonly buffer SourceBuf { uint source[]; };
layout(set = 0, binding = 3, std430) writeonly buffer CompactedBuf { uint compacted[]; };
layout(set = 0, binding = 4, std430) buffer CountBuf { uint count[]; };
layout(set = 0, binding = 5, std140) uniform Params {
uint N; // total number of elements
uint _pad0;
uint _pad1;
uint _pad2;
};
void main() {
uint gid = gl_GlobalInvocationID.x;
if (gid >= N) return;
if (flags[gid] != 0u) {
uint dest = scan_result[gid];
compacted[dest] = source[gid];
}
// Last thread writes total count.
if (gid == N - 1) {
count[0] = scan_result[gid] + flags[gid];
}
}