1#import bevy_hanabi::vfx_common::{
2    ChildInfo, ChildInfoBuffer, EventBuffer, IndirectDispatch, IndirectBuffer,
3    EffectMetadata, RenderGroupIndirect, SimParams, Spawner,
4    seed, tau, pcg_hash, to_float01, frand, frand2, frand3, frand4,
5    rand_uniform_f, rand_uniform_vec2, rand_uniform_vec3, rand_uniform_vec4,
6    rand_normal_f, rand_normal_vec2, rand_normal_vec3, rand_normal_vec4, proj
7}
8
9struct Particle {
10{{ATTRIBUTES}}
11}
12
13struct ParticleBuffer {
14    particles: array<Particle>,
15}
16
17#ifdef READ_PARENT_PARTICLE
18
19struct ParentParticle {
20    {{PARENT_ATTRIBUTES}}
21}
22
23struct ParentParticleBuffer {
24    particles: array<ParentParticle>,
25}
26
27#endif
28
29{{PROPERTIES}}
30
31@group(0) @binding(0) var<uniform> sim_params : SimParams;
32
33// "particle" group @1
34@group(1) @binding(0) var<storage, read_write> particle_buffer : ParticleBuffer;
35@group(1) @binding(1) var<storage, read_write> indirect_buffer : IndirectBuffer;
36#ifdef READ_PARENT_PARTICLE
37@group(1) @binding(2) var<storage, read> parent_particle_buffer : ParentParticleBuffer;
38#endif
39
40// "spawner" group @2
41@group(2) @binding(0) var<storage, read> spawner : Spawner;
42{{PROPERTIES_BINDING}}
43
44// "metadata" group @3
45@group(3) @binding(0) var<storage, read_write> effect_metadata : EffectMetadata;
46#ifdef EMITS_GPU_SPAWN_EVENTS
47{{EMIT_EVENT_BUFFER_BINDINGS}}
48#endif
49
50{{UPDATE_EXTRA}}
51
52#ifdef EMITS_GPU_SPAWN_EVENTS
53{{EMIT_EVENT_BUFFER_APPEND_FUNCS}}
54#endif
55
56@compute @workgroup_size(64)
57fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
58    let thread_index = global_invocation_id.x;
59
60    // Cap at maximum number of alive particles.
61    if (thread_index >= effect_metadata.max_update) {
62        return;
63    }
64
65    // Always write into ping, read from pong
66    let write_index = effect_metadata.ping;
67    let read_index = 1u - write_index;
68
69    let particle_index = indirect_buffer.indices[3u * thread_index + read_index];
70
71    // Initialize the PRNG seed
72    seed = pcg_hash(particle_index ^ spawner.seed);
73
74    var particle: Particle = particle_buffer.particles[particle_index];
75    {{AGE_CODE}}
76    {{REAP_CODE}}
77    {{UPDATE_CODE}}
78
79    {{WRITEBACK_CODE}}
80
81    // Check if alive
82    if (!is_alive) {
83        // Save dead index
84        let dead_index = atomicAdd(&effect_metadata.dead_count, 1u);
85        indirect_buffer.indices[3u * dead_index + 2u] = particle_index;
86
87        // Also increment copy of dead count, which was updated in dispatch indirect
88        // pass just before, and need to remain correct after this pass
89        atomicAdd(&effect_metadata.max_spawn, 1u);
90        atomicSub(&effect_metadata.alive_count, 1u);
91    } else {
92        // Increment alive particle count and write indirection index for later rendering
93        let indirect_index = atomicAdd(&effect_metadata.instance_count, 1u);
94        indirect_buffer.indices[3u * indirect_index + write_index] = particle_index;
95    }
96}