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.instance_count); // TODO - atomic not needed
36    if (thread_index >= count) {
37        return;
38    }
39
40    let read_index = effect_metadata.ping;
41    let particle_index = indirect_index_buffer[thread_index * 3u + read_index];
42
43    let particle_offset = particle_index * effect_metadata.particle_stride;
44    let key_offset = particle_offset + effect_metadata.sort_key_offset;
45    let key2_offset = particle_offset + effect_metadata.sort_key2_offset;
46
47    let pair_index = atomicAdd(&sort_buffer.count, 1);
48    sort_buffer.pairs[pair_index].key = particle_buffer.data[key_offset];
49    sort_buffer.pairs[pair_index].key2 = particle_buffer.data[key2_offset];
50    sort_buffer.pairs[pair_index].value = particle_index;
51}