1#import bevy_hanabi::vfx_common::{
2    ChildInfo, ChildInfoBuffer, EventBuffer, DispatchIndirectArgs, IndirectBuffer,
3    EffectMetadata, RenderGroupIndirect, SimParams, Spawner, DrawIndexedIndirectArgs,
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@group(0) @binding(1) var<storage, read_write> draw_indirect_buffer : array<DrawIndexedIndirectArgs>;
33
34// "particle" group @1
35@group(1) @binding(0) var<storage, read_write> particle_buffer : ParticleBuffer;
36@group(1) @binding(1) var<storage, read_write> indirect_buffer : IndirectBuffer;
37#ifdef READ_PARENT_PARTICLE
38@group(1) @binding(2) var<storage, read> parent_particle_buffer : ParentParticleBuffer;
39#endif
40
41// "spawner" group @2
42@group(2) @binding(0) var<storage, read> spawner : Spawner;
43{{PROPERTIES_BINDING}}
44
45// "metadata" group @3
46@group(3) @binding(0) var<storage, read_write> effect_metadatas : array<EffectMetadata>;
47#ifdef EMITS_GPU_SPAWN_EVENTS
48{{EMIT_EVENT_BUFFER_BINDINGS}}
49#endif
50
51{{UPDATE_EXTRA}}
52
53#ifdef EMITS_GPU_SPAWN_EVENTS
54{{EMIT_EVENT_BUFFER_APPEND_FUNCS}}
55#endif
56
57@compute @workgroup_size(64)
58fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
59    let thread_index = global_invocation_id.x;
60
61    // Cap at maximum number of alive particles.
62    let effect_metadata = &effect_metadatas[spawner.effect_metadata_index];
63    if (thread_index >= (*effect_metadata).max_update) {
64        return;
65    }
66
67    let base_particle = spawner.slab_offset;
68#ifdef READ_PARENT_PARTICLE
69    let parent_base_particle = spawner.parent_slab_offset;
70#endif
71
72    let write_index = effect_metadata.indirect_write_index;
73    let read_index = 1u - write_index;
74
75    let particle_index = indirect_buffer
76        .rows[base_particle + thread_index]
77        .particle_index[read_index];
78
79    // Initialize the PRNG seed
80    seed = pcg_hash(particle_index ^ spawner.seed);
81
82    var particle: Particle = particle_buffer.particles[base_particle + particle_index];
83    {{AGE_CODE}}
84    {{REAP_CODE}}
85    {{UPDATE_CODE}}
86
87    {{WRITEBACK_CODE}}
88
89    // Check if alive
90    if (!is_alive) {
91        // Save dead index
92        let alive_index = atomicSub(&((*effect_metadata).alive_count), 1u) - 1u;
93        indirect_buffer.rows[base_particle + alive_index].dead_index = base_particle + particle_index;
94
95        // DEBUG
96        //indirect_buffer.rows[base_particle + alive_index].particle_index[0] = 0xFFFFFFFFu;
97        //indirect_buffer.rows[base_particle + alive_index].particle_index[1] = 0xFFFFFFFFu;
98
99        // Also increment copy of dead count, which was updated in dispatch indirect
100        // pass just before, and need to remain correct after this pass. We wouldn't have
101        // to do that here if we had a per-effect pass between update and the next init.
102        atomicAdd(&((*effect_metadata).max_spawn), 1u);
103    } else {
104        // Increment visible particle count (in the absence of any GPU culling), and write
105        // the indirection index for later rendering.
106        let indirect_index = atomicAdd(&draw_indirect_buffer[(*effect_metadata).indirect_render_index].instance_count, 1u);
107        indirect_buffer.rows[base_particle + indirect_index].particle_index[write_index] = particle_index;
108    }
109}