1#import bevy_hanabi::vfx_common::{EffectMetadata, Spawner}
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@group(0) @binding(3) var<storage, read> spawner : Spawner;
28
29/// Copy the sorted particle indices back into the effect index buffer.
30@compute @workgroup_size(64)
31fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
32    let row_index = global_invocation_id.x;
33    let count = atomicLoad(&effect_metadata.alive_count); // TODO - atomic not needed
34    if (row_index >= count) {
35        return;
36    }
37
38    let base_particle = spawner.slab_offset;
39
40    // Write the particle index back into the same ping-pong column index we read it from
41    // in vfx_sort_fill. Sorting is optional and shouldn't influence that ping-pong logic.
42    let write_index = effect_metadata.indirect_write_index;
43
44    let particle_index = sort_buffer.pairs[row_index].value;
45    indirect_index_buffer.data[(base_particle + row_index) * 3u + write_index] = particle_index;
46}