1#import bevy_hanabi::vfx_common::{
2    ChildInfo, ChildInfoBuffer, EventBuffer, DispatchIndirectArgs, 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_metadatas : array<EffectMetadata>;
46#ifdef CONSUME_GPU_SPAWN_EVENTS
47@group(3) @binding(1) var<storage, read> child_info_buffer : ChildInfoBuffer;
48@group(3) @binding(2) var<storage, read> event_buffer : EventBuffer;
49#endif
50
51{{INIT_EXTRA}}
52
53@compute @workgroup_size(64)
54fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
55    let thread_index = global_invocation_id.x;
56
57    // Cap to max number of dead particles, copied from (capacity - alive_count) at the end
58    // of the previous iteration, and constant during this pass (unlike alive_count).
59    let effect_metadata = &effect_metadatas[spawner.effect_metadata_index];
60    let max_spawn = atomicLoad(&(*effect_metadata).max_spawn);
61    if (thread_index >= max_spawn) {
62        return;
63    }
64
65    // Cap to the actual number of spawning requested by CPU or GPU, since compute shaders run
66    // in workgroup_size(64) so more threads than needed are launched (rounded up to 64).
67#ifdef CONSUME_GPU_SPAWN_EVENTS
68    let event_index = thread_index;
69    let global_child_index = (*effect_metadata).global_child_index;
70    let event_count = child_info_buffer.rows[global_child_index].event_count;
71    if (event_index >= u32(event_count)) {
72        return;
73    }
74#else
75    // Cap to the actual number of spawning requested by CPU (in the case of
76    // spawners) or the number of particles present in the source group (in the
77    // case of cloners).
78    let spawn_count: u32 = u32(spawner.spawn);
79    if (thread_index >= spawn_count) {
80        return;
81    }
82#endif
83
84    let base_particle = spawner.slab_offset;
85
86    // Count as alive, and recycle a dead particle slot to store the newly spawned particle
87    let alive_index = atomicAdd(&(*effect_metadata).alive_count, 1u);
88    let slab_particle_index = indirect_buffer.rows[base_particle + alive_index].dead_index;
89    let particle_index = slab_particle_index - base_particle;
90
91    // DEBUG
92    //indirect_buffer.rows[base_particle + alive_index].dead_index = 0xFFFFFFFFu;
93
94    // Bump the particle counter each time we allocate a particle. This generates a unique
95    // particle ID used for various purposes (but not directly by the simulation). We still
96    // store it in a variable, because the INIT_CODE might access it.
97    let particle_counter = atomicAdd(&(*effect_metadata).particle_counter, 1u);
98
99    // Initialize the PRNG seed
100    seed = pcg_hash(particle_index ^ spawner.seed);
101
102    // Spawner transform
103    let transform = transpose(
104        mat4x4(
105            spawner.transform[0],
106            spawner.transform[1],
107            spawner.transform[2],
108            vec4<f32>(0.0, 0.0, 0.0, 1.0)
109        )
110    );
111
112#ifdef READ_PARENT_PARTICLE
113    // Fetch parent particle which triggered this spawn
114    let parent_base_particle = spawner.parent_slab_offset;
115    let parent_particle_index = event_buffer.spawn_events[event_index].particle_index;
116    let parent_particle = parent_particle_buffer.particles[parent_base_particle + parent_particle_index];
117#endif
118
119    // Initialize new particle
120    var particle = Particle();
121    {{INIT_CODE}}
122#ifdef ATTRIBUTE_PREV
123    particle.prev = 0xffffffffu;
124#endif  // ATTRIBUTE_PREV
125#ifdef ATTRIBUTE_NEXT
126    particle.next = 0xffffffffu;
127#endif  // ATTRIBUTE_NEXT
128
129    // Only add emitter's transform to CPU-spawned particles. GPU-spawned particles
130    // don't have a CPU spawner to read from (they can read the parent particle
131    // transform, which is equivalent).
132#ifndef CONSUME_GPU_SPAWN_EVENTS
133    {{SIMULATION_SPACE_TRANSFORM_PARTICLE}}
134#endif
135
136    // Append to alive list of indirect buffer.
137    let write_index = (*effect_metadata).indirect_write_index;
138    indirect_buffer.rows[base_particle + alive_index].particle_index[write_index] = particle_index;
139
140    // Write back new particle
141    particle_buffer.particles[base_particle + particle_index] = particle;
142}