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: 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@group(0) @binding(4) var<storage, read> spawner : Spawner;
31
32/// Fill the sorting key-value pair buffer with data to prepare for actual sorting.
33@compute @workgroup_size(64)
34fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
35    let thread_index = global_invocation_id.x;
36    let count = atomicLoad(&effect_metadata.alive_count); // TODO - atomic not needed
37    if (thread_index >= count) {
38        return;
39    }
40
41    let base_particle = spawner.slab_offset;
42
43    // Read the particle index from the ping-pong column index written this frame by vfx_update.
44    // After sorting we will write back into that same column, so that the sorting effectively
45    // sorted the column in-place (no ping-pong column swap here).
46    let read_index = effect_metadata.indirect_write_index;
47    let particle_index = indirect_index_buffer[(base_particle + thread_index) * 3u + read_index];
48
49    let particle_offset = (base_particle + particle_index) * effect_metadata.particle_stride;
50    let key_offset = particle_offset + effect_metadata.sort_key_offset;
51    let key2_offset = particle_offset + effect_metadata.sort_key2_offset;
52
53    let pair_index = atomicAdd(&sort_buffer.count, 1);
54    sort_buffer.pairs[pair_index].key = particle_buffer.data[key_offset];
55    sort_buffer.pairs[pair_index].key2 = particle_buffer.data[key2_offset];
56    sort_buffer.pairs[pair_index].value = particle_index;
57}