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: atomic<i32>,
16    pairs: array<KeyValuePair>,
17}
18
19/// Particle buffer as an array of u32. This prevents having to specialize this shader
20/// for every single particle layout.
21struct RawParticleBuffer {
22    data: array<u32>,
23}
24
25@group(0) @binding(0) var<storage, read_write> sort_buffer : SortBuffer;
26@group(0) @binding(1) var<storage, read> particle_buffer : RawParticleBuffer;
27@group(0) @binding(2) var<storage, read> indirect_index_buffer : array<u32>;
28// Technically read-only, but the type contains atomic<> fields and wasm is strict about it
29@group(0) @binding(3) var<storage, read_write> effect_metadata : EffectMetadata;
30
31/// Fill the sorting key-value pair buffer with data to prepare for actual sorting.
32@compute @workgroup_size(64)
33fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
34    let thread_index = global_invocation_id.x;
35    let count = atomicLoad(&effect_metadata.alive_count); // TODO - atomic not needed
36    if (thread_index >= count) {
37        return;
38    }
39
40    // Read the particle index from the ping-pong column index written this frame by vfx_update.
41    // After sorting we will write back into that same column, so that the sorting effectively
42    // sorted the column in-place (no ping-pong column swap here).
43    let read_index = effect_metadata.indirect_write_index;
44    let particle_index = indirect_index_buffer[thread_index * 3u + read_index];
45
46    let particle_offset = particle_index * effect_metadata.particle_stride;
47    let key_offset = particle_offset + effect_metadata.sort_key_offset;
48    let key2_offset = particle_offset + effect_metadata.sort_key2_offset;
49
50    let pair_index = atomicAdd(&sort_buffer.count, 1);
51    sort_buffer.pairs[pair_index].key = particle_buffer.data[key_offset];
52    sort_buffer.pairs[pair_index].key2 = particle_buffer.data[key2_offset];
53    sort_buffer.pairs[pair_index].value = particle_index;
54}