1#import bevy_hanabi::vfx_common::{
2    ChildInfo, ChildInfoBuffer, EventBuffer, DispatchIndirectArgs, IndirectBuffer,
3    EffectMetadata, RenderGroupIndirect, SimParams, Spawner, BatchInfo,
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/// Location of an effect in a slab.
32struct EffectLocation {
33    /// Index of the effect in the global list of effects.
34    effect_index: u32,
35    /// Base particle index, that is index in the slab of the first particle for this instance.
36    base_particle: u32,
37    /// Index of this particle relative to its effect. Note that if there's an indirection
38    /// buffer then this is the linear index in [0:N[ of the particle to update, before the indirection.
39    update_index: u32,
40}
41
42/// Find the index of an effect from the index of a particle.
43///
44/// This uses a binary search on the slab_offset field of the spawners array, which
45/// represents a prefix sum of the particle count per effect (for previous effects;
46/// the value is actually the base particle so the first entry is always 0).
47///
48/// Requirements:
49/// - var<storage, read> batch_info : BatchInfo
50/// - var<storage, read> prefix_sum : array<u32>
51fn find_location_from_particle(slab_particle_index: u32) -> EffectLocation {
52    var lo = batch_info.prefix_sum_offset;
53    var hi = lo + batch_info.prefix_sum_count;
54    var num_iter = 0;  // avoid deadlocking the GPU by capping the iteration count
55    while (lo < hi) {
56        let mid = (hi + lo) >> 1u;
57        let base_particle = prefix_sum[mid];
58        if (slab_particle_index >= base_particle) {
59            lo = mid + 1u;
60        } else if (slab_particle_index < base_particle) {
61            hi = mid;
62        }
63        num_iter += 1;
64        if (num_iter >= 100) {
65            return EffectLocation(0xDEADBEEFu, 0xDEADBEEFu, 0xDEADBEEFu);
66        }
67    }
68    let base_particle = prefix_sum[lo - 1u];
69    let effect_index = lo - 1u - batch_info.prefix_sum_offset;
70    let update_index = slab_particle_index - base_particle;
71    return EffectLocation(effect_index, base_particle, update_index);
72}
73
74@group(0) @binding(0) var<uniform> sim_params : SimParams;
75
76// "particle" group @1
77@group(1) @binding(0) var<storage, read_write> particle_buffer : ParticleBuffer;
78@group(1) @binding(1) var<storage, read_write> indirect_buffer : IndirectBuffer;
79#ifdef READ_PARENT_PARTICLE
80@group(1) @binding(2) var<storage, read> parent_particle_buffer : ParentParticleBuffer;
81#endif
82
83// "spawner" group @2
84@group(2) @binding(0) var<storage, read> spawners : array<Spawner>;
85@group(2) @binding(1) var<storage, read> prefix_sum : array<u32>;
86@group(2) @binding(2) var<storage, read> batch_info : BatchInfo;
87{{PROPERTIES_BINDING}}
88
89// "metadata" group @3
90@group(3) @binding(0) var<storage, read_write> effect_metadatas : array<EffectMetadata>;
91#ifdef CONSUME_GPU_SPAWN_EVENTS
92@group(3) @binding(1) var<storage, read> child_info_buffer : ChildInfoBuffer;
93@group(3) @binding(2) var<storage, read> event_buffer : EventBuffer;
94#endif
95
96{{INIT_EXTRA}}
97
98var<private> effect_metadata_index: u32;
99var<private> properties_array_index: u32;
100
101@compute @workgroup_size(64)
102fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
103    // Global particle index into the slab, including those particles from other
104    // effect instances in the same batch, as well as possibly from other batches.
105    // This is rarely useful on its own.
106    let slab_particle_index = batch_info.base_particle + global_invocation_id.x;
107
108    // Find the index of the effect this particle is part of.
109    let location = find_location_from_particle(slab_particle_index);
110    let spawner = &spawners[batch_info.base_effect + location.effect_index];
111    effect_metadata_index = (*spawner).effect_metadata_index;
112    let base_particle = location.base_particle;
113
114    // Cap to max number of dead particles, copied from (capacity - alive_count) at the end
115    // of the previous iteration, and constant during this pass (unlike alive_count).
116    let effect_metadata = &effect_metadatas[(*spawner).effect_metadata_index];
117    let max_spawn = atomicLoad(&(*effect_metadata).max_spawn);
118    if (location.update_index >= max_spawn) {
119        return;
120    }
121    properties_array_index = (*effect_metadata).properties_array_index;
122
123    // Cap to the actual number of spawning requested by CPU or GPU, since compute shaders run
124    // in workgroup_size(64) so more threads than needed are launched (rounded up to 64).
125#ifdef CONSUME_GPU_SPAWN_EVENTS
126    let event_index = location.update_index;
127    let global_child_index = (*effect_metadata).global_child_index;
128    let event_count = child_info_buffer.rows[global_child_index].event_count;
129    if (event_index >= u32(event_count)) {
130        return;
131    }
132#else
133    // Cap to the actual number of spawning requested by CPU (in the case of
134    // spawners) or the number of particles present in the source group (in the
135    // case of cloners).
136    let spawn_count: u32 = u32((*spawner).spawn);
137    if (location.update_index >= spawn_count) {
138        return;
139    }
140#endif
141
142    // Count as alive, and recycle a dead particle slot to store the newly spawned particle
143    let alive_index = atomicAdd(&(*effect_metadata).alive_count, 1u);
144    let slab_particle_dead_index = indirect_buffer.rows[base_particle + alive_index].dead_index;
145    let particle_index = slab_particle_dead_index - base_particle;
146
147    // DEBUG
148    //indirect_buffer.rows[base_particle + alive_index].dead_index = 0xFFFFFFFFu;
149
150    // Bump the particle counter each time we allocate a particle. This generates a unique
151    // particle ID used for various purposes (but not directly by the simulation). We still
152    // store it in a variable, because the INIT_CODE might access it.
153    let particle_counter = atomicAdd(&(*effect_metadata).particle_counter, 1u);
154
155    // Initialize the PRNG seed
156    seed = pcg_hash(particle_index ^ (*spawner).seed);
157
158    // Spawner transform
159    let transform = transpose(
160        mat4x4(
161            (*spawner).transform[0],
162            (*spawner).transform[1],
163            (*spawner).transform[2],
164            vec4<f32>(0.0, 0.0, 0.0, 1.0)
165        )
166    );
167
168#ifdef READ_PARENT_PARTICLE
169    // Fetch parent particle which triggered this spawn
170    let parent_base_particle = (*spawner).parent_slab_offset;
171    let parent_particle_index = event_buffer.spawn_events[event_index].particle_index;
172    let parent_particle = parent_particle_buffer.particles[parent_base_particle + parent_particle_index];
173#endif
174
175    // Initialize new particle
176    var particle = Particle();
177    {{INIT_CODE}}
178#ifdef ATTRIBUTE_PREV
179    particle.prev = 0xffffffffu;
180#endif  // ATTRIBUTE_PREV
181#ifdef ATTRIBUTE_NEXT
182    particle.next = 0xffffffffu;
183#endif  // ATTRIBUTE_NEXT
184
185    // Only add emitter's transform to CPU-spawned particles. GPU-spawned particles
186    // don't have a CPU spawner to read from (they can read the parent particle
187    // transform, which is equivalent).
188#ifndef CONSUME_GPU_SPAWN_EVENTS
189    {{SIMULATION_SPACE_TRANSFORM_PARTICLE}}
190#endif
191
192    // Append to alive list of indirect buffer.
193    let write_index = (*effect_metadata).indirect_write_index;
194    indirect_buffer.rows[base_particle + alive_index].particle_index[write_index] = particle_index;
195
196    // Write back new particle
197    particle_buffer.particles[base_particle + particle_index] = particle;
198}