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@group(2) @binding(1) var<storage, read_write> prefix_sum : array<u32>;
21
22#ifdef HAS_GPU_SPAWN_EVENTS
23@group(3) @binding(0) var<storage, read_write> child_info_buffer : ChildInfoBuffer;
24#endif
25
26/// Perform per-instance update between the init and update passes.
27///
28/// This is invoked once per instance, with the global invocation ID indexing the list of spawners,
29/// which is re-uploaded each frame so is guaranteed to be a tightly packed array with exactly the
30/// number of 
31@compute @workgroup_size(64)
32fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
33    // Cap at maximum number of effects
34    let global_effect_index = global_invocation_id.x;
35    if (global_effect_index >= sim_params.num_effects) {
36        return;
37    }
38
39#ifdef HAS_GPU_SPAWN_EVENTS
40    // Clear any GPU event. The indexing is inconsistent here (indexing child info buffer with an
41    // effect instance), but is safe because there are always less child effects than there are effects
42    // in total, so the index will always cover the entire child info array. And we're only resetting
43    // a value so don't read anything, so only need to visit each entry once but the order doesn't matter.
44    if (global_effect_index < arrayLength(&child_info_buffer.rows)) {
45        child_info_buffer.rows[global_effect_index].event_count = 0;
46    }
47#endif
48
49    let spawner = &spawner_buffer[global_effect_index];
50
51    let effect_metadata_index = (*spawner).effect_metadata_index;
52    let em_base = EFFECT_METADATA_STRIDE * effect_metadata_index;
53
54    // Clear the rendering instance count, which will be upgraded by the update pass
55    // with the particles actually alive at the end of their update (after aged).
56    let draw_indirect_index = (*spawner).draw_indirect_index;
57    let dri_base = DRAW_INDEXED_INDIRECT_STRIDE * draw_indirect_index;
58    draw_indirect_buffer[dri_base + 1u] = 0u;
59
60    let capacity = effect_metadata_buffer[em_base + EM_OFFSET_CAPACITY];
61    let alive_count = effect_metadata_buffer[em_base + EM_OFFSET_ALIVE_COUNT];
62    let dead_count = capacity - alive_count;
63
64    // Prepare to rebuild the prefix sum of active instances by storing in the prefix sum array
65    // the actual number of alive particles for each instance. We will build the prefix sum in
66    // the next pass (vfx_prefix_sum) once that's done.
67    // The prefix sums and spawners are allocated in sync, so are indexed with the same index value.
68    prefix_sum[global_effect_index] = alive_count;
69
70    // Update max_update from current value of alive_count, so that the
71    // update pass coming next can cap its threads to this value, while also
72    // atomically modifying alive_count itself for next frame.
73    effect_metadata_buffer[em_base + EM_OFFSET_MAX_UPDATE] = alive_count;
74
75    // Copy the number of dead particles to a constant location, so that the
76    // init pass on next frame can atomically modify alive_count in parallel
77    // yet still read its initial value at the beginning of the init pass,
78    // and limit the number of particles spawned to the number of dead
79    // particles to recycle.
80    effect_metadata_buffer[em_base + EM_OFFSET_MAX_SPAWN] = dead_count;
81
82    // Calculate the number of workgroups (thread groups) to dispatch for the update
83    // pass, which is the number of alive particles rounded up to 64 (workgroup_size).
84    let indirect_dispatch_index = effect_metadata_buffer[em_base + EM_OFFSET_INDIRECT_DISPATCH_INDEX];
85    dispatch_indirect_buffer[indirect_dispatch_index].x = (alive_count + 63u) >> 6u;
86    dispatch_indirect_buffer[indirect_dispatch_index].y = 1u;
87    dispatch_indirect_buffer[indirect_dispatch_index].z = 1u;
88
89    // Swap ping/pong buffers. The update pass always writes into ping, and both the update
90    // pass and the render pass always read from pong.
91    let ping = effect_metadata_buffer[em_base + EM_OFFSET_INDIRECT_WRITE_INDEX];
92    let pong = 1u - ping;
93    effect_metadata_buffer[em_base + EM_OFFSET_INDIRECT_WRITE_INDEX] = pong;
94
95    // Copy the new pong into the spawner buffer, which will be used during rendering
96    // to determine where to read particle indices.
97    (*spawner).render_indirect_read_index = pong;
98}