1#import bevy_hanabi::vfx_common::EffectMetadata
2
3/// Key-value pair for sorting, with optional second sort key.
4struct KeyValuePair {
5    /// Sorting key.
6    key: u32,
7    /// Secondary sorting key. Sorts value with the same primary key.
8    key2: u32,
9    /// Value associated with the sort key(s), generally an index to some other data.
10    /// Copied as is and otherwise ignored by the sorting algorithm.
11    value: u32,
12}
13
14struct SortBuffer {
15    count: i32,
16    pairs: array<KeyValuePair>,
17}
18
19struct IndirectIndexBuffer {
20    data: array<u32>,
21}
22
23@group(0) @binding(0) var<storage, read_write> indirect_index_buffer : IndirectIndexBuffer;
24@group(0) @binding(1) var<storage, read> sort_buffer : SortBuffer;
25// Technically read-only, but the type contains atomic<> fields and wasm is strict about it
26@group(0) @binding(2) var<storage, read_write> effect_metadata : EffectMetadata;
27
28/// Copy the sorted particle indices back into the effect index buffer.
29@compute @workgroup_size(64)
30fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
31    let row_index = global_invocation_id.x;
32    let count = atomicLoad(&effect_metadata.alive_count); // TODO - atomic not needed
33    if (row_index >= count) {
34        return;
35    }
36
37    // Write the particle index back into the same ping-pong column index we read it from
38    // in vfx_sort_fill. Sorting is optional and shouldn't influence that ping-pong logic.
39    let write_index = effect_metadata.indirect_write_index;
40
41    let particle_index = sort_buffer.pairs[row_index].value;
42    indirect_index_buffer.data[row_index * 3u + write_index] = particle_index;
43}