1#import bevy_hanabi::vfx_common::{
2    ChildInfo, ChildInfoBuffer, EventBuffer, DispatchIndirectArgs, IndirectBuffer,
3    EffectMetadata, RenderGroupIndirect, SimParams, Spawner, DrawIndexedIndirectArgs, 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 = batch_info.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@group(0) @binding(1) var<storage, read_write> draw_indirect_buffer : array<DrawIndexedIndirectArgs>;
76
77// "particle" group @1
78@group(1) @binding(0) var<storage, read_write> particle_buffer : ParticleBuffer;
79@group(1) @binding(1) var<storage, read_write> indirect_buffer : IndirectBuffer;
80#ifdef READ_PARENT_PARTICLE
81@group(1) @binding(2) var<storage, read> parent_particle_buffer : ParentParticleBuffer;
82#endif
83
84// "spawner" group @2
85@group(2) @binding(0) var<storage, read> spawners : array<Spawner>;
86@group(2) @binding(1) var<storage, read> prefix_sum : array<u32>;
87@group(2) @binding(2) var<storage, read> batch_info : BatchInfo;
88{{PROPERTIES_BINDING}}
89
90// "metadata" group @3
91@group(3) @binding(0) var<storage, read_write> effect_metadatas : array<EffectMetadata>;
92#ifdef EMITS_GPU_SPAWN_EVENTS
93{{EMIT_EVENT_BUFFER_BINDINGS}}
94#endif
95
96{{UPDATE_EXTRA}}
97
98#ifdef EMITS_GPU_SPAWN_EVENTS
99{{EMIT_EVENT_BUFFER_APPEND_FUNCS}}
100#endif
101
102var<private> effect_metadata_index: u32;
103var<private> properties_array_index: u32;
104
105@compute @workgroup_size(64)
106fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
107    // Global particle index into the slab, including those particles from other
108    // effect instances in the same batch, as well as possibly from other batches.
109    // This is rarely useful on its own.
110    let slab_particle_index = batch_info.base_particle + global_invocation_id.x;
111
112    // Find the index of the effect this particle is part of.
113    let location = find_location_from_particle(slab_particle_index);
114    let spawner = &spawners[batch_info.base_effect + location.effect_index];
115    effect_metadata_index = (*spawner).effect_metadata_index;
116    let base_particle = location.base_particle;
117
118    // Cap at maximum number of alive particles for the current effect
119    let effect_metadata = &effect_metadatas[effect_metadata_index];
120    if (location.update_index >= (*effect_metadata).max_update) {
121        return;
122    }
123    properties_array_index = (*effect_metadata).properties_array_index;
124
125    let write_index = effect_metadata.indirect_write_index;
126    let read_index = 1u - write_index;
127
128    // This is the actual particle index, from the indirection buffer, addressing an
129    // actually alive particle.
130    let particle_index = indirect_buffer
131        .rows[slab_particle_index]
132        .particle_index[read_index];
133
134#ifdef READ_PARENT_PARTICLE
135    let parent_base_particle = (*spawner).parent_slab_offset;
136#endif
137
138    // Initialize the PRNG seed
139    seed = pcg_hash(particle_index ^ (*spawner).seed);
140
141    var particle: Particle = particle_buffer.particles[base_particle + particle_index];
142    {{AGE_CODE}}
143    {{REAP_CODE}}
144    {{UPDATE_CODE}}
145
146    {{WRITEBACK_CODE}}
147
148    // Check if alive
149    if (!is_alive) {
150        // Save dead index. Note that dead_index is a global slab index.
151        let alive_index = atomicSub(&((*effect_metadata).alive_count), 1u) - 1u;
152        indirect_buffer.rows[base_particle + alive_index].dead_index = base_particle + particle_index;
153
154        // DEBUG
155        //indirect_buffer.rows[base_particle + alive_index].particle_index[0] = 0xFFFFFFFFu;
156        //indirect_buffer.rows[base_particle + alive_index].particle_index[1] = 0xFFFFFFFFu;
157
158        // Also increment copy of dead count, which was updated in dispatch indirect
159        // pass just before, and need to remain correct after this pass. We wouldn't have
160        // to do that here if we had a per-effect pass between update and the next init.
161        atomicAdd(&((*effect_metadata).max_spawn), 1u);
162    } else {
163        // Increment visible particle count (in the absence of any GPU culling), and write
164        // the indirection index for later rendering.
165        let indirect_index = atomicAdd(&draw_indirect_buffer[(*effect_metadata).indirect_render_index].instance_count, 1u);
166        indirect_buffer.rows[base_particle + indirect_index].particle_index[write_index] = particle_index;
167    }
168}