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