1#import bevy_hanabi::vfx_common::{
2    ChildInfo, ChildInfoBuffer, SimParams, Spawner,
3    EM_OFFSET_ALIVE_COUNT, EM_OFFSET_MAX_UPDATE, EM_OFFSET_DEAD_COUNT,
4    EM_OFFSET_MAX_SPAWN, EM_OFFSET_INSTANCE_COUNT, EM_OFFSET_INDIRECT_DISPATCH_INDEX,
5    EM_OFFSET_PING, DISPATCH_INDIRECT_STRIDE, EFFECT_METADATA_STRIDE
6}
7
8@group(0) @binding(0) var<uniform> sim_params : SimParams;
9
10// Tightly packed array of EffectMetadata[], accessed as u32 array.
11@group(1) @binding(0) var<storage, read_write> effect_metadata_buffer : array<u32>;
12// Tightly packed array of IndirectDispatch[], accessed as u32 array.
13@group(1) @binding(1) var<storage, read_write> dispatch_indirect_buffer : array<u32>;
14
15@group(2) @binding(0) var<storage, read_write> spawner_buffer : array<Spawner>;
16
17#ifdef HAS_GPU_SPAWN_EVENTS
18@group(3) @binding(0) var<storage, read_write> child_info_buffer : ChildInfoBuffer;
19#endif
20
21/// Calculate the indirect workgroups counts based on the number of particles alive.
22@compute @workgroup_size(64)
23fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
24    let thread_index = global_invocation_id.x;
25
26#ifdef HAS_GPU_SPAWN_EVENTS
27    // Clear any GPU event. The indexing is safe because there are always less child effects
28    // than there are effects in total, so 'index' will always cover the entire child info array.
29    if (thread_index < arrayLength(&child_info_buffer.rows)) {
30        child_info_buffer.rows[thread_index].event_count = 0;
31    }
32#endif
33
34    // Cap at maximum number of effects
35    let effect_index = thread_index;
36    if (effect_index >= sim_params.num_effects) {
37        return;
38    }
39
40    let effect_metadata_index = spawner_buffer[effect_index].effect_metadata_index;
41    let em_base = EFFECT_METADATA_STRIDE * effect_metadata_index;
42
43    // Clear the rendering instance count, which will be upgraded by the update pass
44    // with the particles actually alive at the end of their update (after aged).
45    effect_metadata_buffer[em_base + EM_OFFSET_INSTANCE_COUNT] = 0u;
46
47    let alive_count = effect_metadata_buffer[em_base + EM_OFFSET_ALIVE_COUNT];
48    let dead_count = effect_metadata_buffer[em_base + EM_OFFSET_DEAD_COUNT];
49
50    // Update max_update from current value of alive_count, so that the
51    // update pass coming next can cap its threads to this value, while also
52    // atomically modifying alive_count itself for next frame.
53    effect_metadata_buffer[em_base + EM_OFFSET_MAX_UPDATE] = alive_count;
54
55    // Copy the number of dead particles to a constant location, so that the
56    // init pass on next frame can atomically modify dead_count in parallel
57    // yet still read its initial value at the beginning of the init pass,
58    // and limit the number of particles spawned to the number of dead
59    // particles to recycle.
60    effect_metadata_buffer[em_base + EM_OFFSET_MAX_SPAWN] = dead_count;
61
62    // Calculate the number of workgroups (thread groups) to dispatch for the update
63    // pass, which is the number of alive particles rounded up to 64 (workgroup_size).
64    let indirect_dispatch_index = effect_metadata_buffer[em_base + EM_OFFSET_INDIRECT_DISPATCH_INDEX];
65    let di_base = DISPATCH_INDIRECT_STRIDE * indirect_dispatch_index;
66    dispatch_indirect_buffer[di_base] = (alive_count + 63u) >> 6u;
67    dispatch_indirect_buffer[di_base + 1u] = 1u;
68    dispatch_indirect_buffer[di_base + 2u] = 1u;
69
70    // Swap ping/pong buffers. The update pass always writes into ping, and both the update
71    // pass and the render pass always read from pong.
72    let ping = effect_metadata_buffer[em_base + EM_OFFSET_PING];
73    let pong = 1u - ping;
74    effect_metadata_buffer[em_base + EM_OFFSET_PING] = pong;
75
76    // Copy the new pong into the spawner buffer, which will be used during rendering
77    // to determine where to read particle indices.
78    spawner_buffer[effect_index].render_pong = pong;
79}