1#define_import_path bevy_hanabi::vfx_common
2
3struct SimParams {
4    /// Delta time in seconds since last simulation tick.
5    delta_time: f32,
6    /// Time in seconds since the start of simulation.
7    time: f32,
8    /// Virtual delta time in seconds since last simulation tick.
9    virtual_delta_time: f32,
10    /// Virtual time in seconds since the start of simulation.
11    virtual_time: f32,
12    /// Real delta time in seconds since last simulation tick.
13    real_delta_time: f32,
14    /// Real time in seconds since the start of simulation.
15    real_time: f32,
16    /// Total number of effects to update this frame. Used by the indirect
17    /// compute pipeline to cap the compute thread to the actual number of
18    /// effects to process.
19    num_effects: u32,
20}
21
22struct Spawner {
23    /// Compressed transform of the emitter.
24    transform: mat3x4<f32>, // transposed (row-major)
25    /// Inverse compressed transform of the emitter.
26    inverse_transform: mat3x4<f32>, // transposed (row-major)
27    /// Number of particles to spawn this frame, as calculated by the CPU Spawner.
28    ///
29    /// This is only used if the effect is not a child effect (driven by GPU events).
30    spawn: i32,
31    /// PRNG seed for this effect instance. Currently this can change each time the
32    /// effect is recompiled, and cannot be set deterministically (TODO).
33    seed: u32,
34    /// Index of the ping-pong buffer of particle indices to read particles from
35    /// during rendering. Cached from EffectMetadata::ping after it's swapped
36    /// in the indirect dispatch, because the EffectMetadata struct is used by GPU
37    /// as an indirect draw source so cannot also be bound as regular storage
38    /// buffer for reading.
39    render_indirect_read_index: u32,
40    /// Index of the [`EffectMetadata`] for this effect.
41    effect_metadata_index: u32,
42    /// Index of the [`DrawIndirectArgs`] or [`DrawIndexedIndirectArgs`] for this effect.
43    draw_indirect_index: u32,
44    /// Start offset of the particles and indirect indices into the effect's
45    /// slab, in number of particles (row index).
46    slab_offset: u32,
47    /// Start offset of the particles and indirect indices into the parent effect's
48    /// slab (if the effect has a parent effect), in number of particles (row index).
49    /// This is ignored if the effect has no parent.
50    parent_slab_offset: u32,
51
52    // Keep this field here separate from auto-padding below, so that the WGSL and Rust
53    // struct definitions match.
54    unused: u32,
55
56    {{SPAWNER_PADDING}}
57}
58
59const SPAWNER_OFFSET_PONG: u32 = 27u;
60
61/// Single row entry into an IndirectBuffer.
62///
63/// Each row corresponds to one row in the ParticleBuffer. The total number of rows in the
64/// IndirectBuffer is the capacity of the ParticleBuffer itself (that is, they have the same
65/// number of rows).
66struct IndirectEntry {
67    /// Ping-pong index into ParticleBuffer.
68    particle_index: array<u32, 2>,
69    /// Index into ParticleBuffer of a dead particle slot which can be recycled.
70    dead_index: u32,
71}
72
73/// Indirection buffer storing the indices of alive and dead particles as a contiguous
74/// range, to ensure we can dispatch a tight number of workgroups only for alive particles,
75/// and can find a dead one to recycle quickly.
76struct IndirectBuffer {
77    rows: array<IndirectEntry>,
78}
79
80/// A event emitted by another effect during its update pass, to trigger the spawning
81/// of one or more particle in this effect.
82struct SpawnEvent {
83    /// The particle index in the parent effect buffer of the source particle which
84    /// triggered the event. This is used to inherit attributes like position or velocity.
85    particle_index: u32,
86}
87
88/// Append buffer populated during the Update pass of the previous frame by a parent effect,
89/// and read back by its child effect(s) during the Init pass of the next frame.
90struct EventBuffer {
91    /// The spawn events themselves.
92    spawn_events: array<SpawnEvent>,
93}
94
95/// Info about a single child of a parent effect.
96struct ChildInfo {
97    /// Index of the effect's DispatchIndirectArgs entry in the global init indirect dispatch array.
98    init_indirect_dispatch_index: u32,
99    /// Number of events in the associated event buffer.
100#ifdef CHILD_INFO_EVENT_COUNT_IS_ATOMIC
101    event_count: atomic<i32>,
102#else
103    event_count: i32,
104#endif
105}
106
107/// Buffer storing all the ChildInfo structs for several effects.
108struct ChildInfoBuffer {
109    /// The child info structs themselves.
110    rows: array<ChildInfo>,
111}
112
113/// Indirect compute dispatch struct for GPU-driven passes. The layout of this struct is dictated by WGSL.
114struct DispatchIndirectArgs {
115    /// Number of workgroups. Each workgroup has exactly 64 threads.
116    x: u32,
117    /// Unused; always 1.
118    y: u32,
119    /// Unused; always 1.
120    z: u32,
121}
122
123/// Indirect draw (non-indexed) dispatch struct for GPU-driven passes. The layout of this struct is dictated by WGSL.
124/// See https://docs.rs/wgpu/latest/wgpu/util/struct.DrawIndirectArgs.html.
125struct DrawIndirectArgs {
126    vertex_count: u32,
127    instance_count: atomic<u32>,
128    first_vertex: u32,
129    first_instance: u32,
130}
131
132/// Stride in u32 count (4 bytes) of the DrawIndirectArgs struct.
133const DRAW_INDIRECT_STRIDE: u32 = 4u;
134
135/// Indirect draw (non-indexed) dispatch struct for GPU-driven passes. The layout of this struct is dictated by WGSL.
136/// See https://docs.rs/wgpu/latest/wgpu/util/struct.DrawIndexedIndirectArgs.html.
137struct DrawIndexedIndirectArgs {
138    index_count: u32,
139    instance_count: atomic<u32>,
140    first_index: u32,
141    base_vertex: i32,
142    first_instance: u32,
143}
144
145/// Stride in u32 count (4 bytes) of the DrawIndexedIndirectArgs struct.
146const DRAW_INDEXED_INDIRECT_STRIDE: u32 = 5u;
147
148/// Shared info for a single batch (single shader invocation).
149struct BatchInfo {
150    total_spawn_count: u32,
151    total_update_count: u32,
152    base_effect: u32,
153    /// Offset to apply to the workgroup thread index to determine the global
154    /// particle index in the currently bound slab. This is often (and ideally)
155    /// zero, but may be > 0 if the entire slab cannot be processed with a
156    /// single invocation.
157    base_particle: u32,
158    /// Offset into the prefix sum buffer of the sum for the first effect of
159    /// this batch.
160    prefix_sum_offset: u32,
161    /// Number of active effects in the batch, and therefore number of
162    /// consecutive prefix sum values for effects of this batch.
163    prefix_sum_count: u32,
164
165    /// Padding for storage buffer alignment. This struct is sometimes bound as part
166    /// of an array, or sometimes individually as a single unit. In the later case,
167    /// we need it to be aligned to the GPU limits of the device. That limit is only
168    /// known at runtime when initializing the WebGPU device.
169    {{BATCH_INFO_PADDING}}
170}
171
172// Effect metadata offsets. Used when accessing a tightly packed array of EffectMetadata
173// as a raw array<u32>, so that we can avoid WGSL struct padding and keep data more compact
174// in the GPU buffer. Each offset corresponds to a field in the EffectMetadata struct.
175// Note that all fields are 4 bytes, so we can index by "number of 4-byte field".
176const EM_OFFSET_CAPACITY: u32 = 0u;
177const EM_OFFSET_ALIVE_COUNT: u32 = 1u;
178const EM_OFFSET_MAX_UPDATE: u32 = 2u;
179const EM_OFFSET_MAX_SPAWN: u32 = 3u;
180const EM_OFFSET_INDIRECT_WRITE_INDEX: u32 = 4u;
181const EM_OFFSET_INDIRECT_DISPATCH_INDEX: u32 = 5u;
182
183/// Metadata describing a single effect instance.
184///
185/// The metadata describes various effect settings, as well we the location in
186/// GPU buffers of the piece of data associated with that instance.
187struct EffectMetadata {
188    /// Total number of particles for this effect. This is the capacity of the effect,
189    /// in number of particles, which is used to sub-allocate various storages for this
190    /// effect, notably in the ParticleBuffer and the IndirectBuffer. This is constant
191    /// for the duration of the effect instance life.
192    capacity: u32,
193
194    /// Number of particles alive. Note that because particles can be simulated even
195    /// when off-screen, in theory this could be greater than instance_count. Currently
196    /// we don't have GPU culling, so in practice this remains strictly equal. But we
197    /// store it separately 1) because this could change in the future, and 2) because
198    /// the indirect render fields above should really be in their own buffer, not here.
199    alive_count: atomic<u32>,
200    /// Maximum number of update threads to run. This is cached from `alive_count`
201    /// during the indirect dispatch, so that the update compute pass can cap its
202    /// thread count while also modifying the actual `alive_count` if some particle
203    /// dies during the update pass.
204    max_update: u32,
205    /// Maxmimum number of init threads to run on next frame. This is cached from
206    /// `capacity - alive_count` during the indirect dispatch of the previous frame,
207    /// so that the init compute pass can cap its thread count while also decrementing
208    /// the actual dead count (increment the `alive_count`) as particles are spawned.
209    max_spawn: atomic<u32>,
210
211    /// Write index into the ping-pong buffer for particle indices. The buffers
212    /// are swapped during the indirect dispatch (although the render pass still uses
213    /// the complement of this to write, so technically ignores that swap).
214    indirect_write_index: u32,
215    /// Index of the [`GpuRenderIndirect`] struct inside the global
216    /// [`EffectsMeta::render_group_dispatch_buffer`].
217    indirect_render_index: u32,
218    /// Offset (in u32 count) of the init indirect dispatch struct inside its
219    /// buffer. This avoids having to align those 16-byte structs to the GPU
220    /// alignment (at least 32 bytes, even 256 bytes on some).
221    init_indirect_dispatch_index: u32,
222    /// Index inside the global array of the spawner struct for this effect instance.
223    //spawner_index: u32,
224    /// Offset (in u32 count) of the start of the property block for this
225    /// effect. This is ignored if the effect doesn't use properties.
226    properties_array_index: u32,
227    /// Index of this effect into its parent's ChildInfo array
228    /// ([`EffectChildren::effect_cache_ids`] and its associated GPU
229    /// array). This starts at zero for the first child of each effect, and is
230    /// only unique per parent, not globally. Only available if this effect is a
231    /// child of another effect (i.e. if it has a parent).
232    local_child_index: u32,
233    /// For children, global index of the ChildInfo into the shared array.
234    global_child_index: u32,
235    /// For parents, base index of the their first ChildInfo into the shared array.
236    base_child_index: u32,
237
238    /// Particle stride, in number of u32.
239    particle_stride: u32,
240    /// Offset from the particle start to the first sort key, in number of u32.
241    sort_key_offset: u32,
242    /// Offset from the particle start to the second sort key, in number of u32.
243    sort_key2_offset: u32,
244
245    /// Atomic counter incremented each time a particle spawns. Useful for
246    /// things like RIBBON_ID or any other use where a unique value is needed.
247    /// The value loops back after some time, but unless some particle lives
248    /// forever there's little chance of repetition.
249    particle_counter: atomic<u32>,
250
251    /// Padding for storage buffer alignment. This struct is sometimes bound as part
252    /// of an array, or sometimes individually as a single unit. In the later case,
253    /// we need it to be aligned to the GPU limits of the device. That limit is only
254    /// known at runtime when initializing the WebGPU device.
255    {{EFFECT_METADATA_PADDING}}
256}
257
258/// Stride, in u32 count, between elements of an array<EffectMetadata>.
259const EFFECT_METADATA_STRIDE: u32 = {{EFFECT_METADATA_STRIDE}} / 4u;
260
261var<private> seed : u32 = 0u;
262
263const tau: f32 = 6.283185307179586476925286766559;
264
265// Rand: PCG
266// https://www.reedbeta.com/blog/hash-functions-for-gpu-rendering/
267fn pcg_hash(input: u32) -> u32 {
268    var state: u32 = input * 747796405u + 2891336453u;
269    var word: u32 = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u;
270    return (word >> 22u) ^ word;
271}
272
273fn to_float01(u: u32) -> f32 {
274    // Note: could generate only 24 bits of randomness
275    return bitcast<f32>((u & 0x007fffffu) | 0x3f800000u) - 1.;
276}
277
278// Random floating-point number in [0:1]
279fn frand() -> f32 {
280    seed = pcg_hash(seed);
281    return to_float01(pcg_hash(seed));
282}
283
284// Random floating-point number in [0:1]^2
285fn frand2() -> vec2<f32> {
286    seed = pcg_hash(seed);
287    var x = to_float01(seed);
288    seed = pcg_hash(seed);
289    var y = to_float01(seed);
290    return vec2<f32>(x, y);
291}
292
293// Random floating-point number in [0:1]^3
294fn frand3() -> vec3<f32> {
295    seed = pcg_hash(seed);
296    var x = to_float01(seed);
297    seed = pcg_hash(seed);
298    var y = to_float01(seed);
299    seed = pcg_hash(seed);
300    var z = to_float01(seed);
301    return vec3<f32>(x, y, z);
302}
303
304// Random floating-point number in [0:1]^4
305fn frand4() -> vec4<f32> {
306    // Each rand() produces 32 bits, and we need 24 bits per component,
307    // so can get away with only 3 calls.
308    var r0 = pcg_hash(seed);
309    var r1 = pcg_hash(r0);
310    var r2 = pcg_hash(r1);
311    seed = r2;
312    var x = to_float01(r0);
313    var r01 = (r0 & 0xff000000u) >> 8u | (r1 & 0x0000ffffu);
314    var y = to_float01(r01);
315    var r12 = (r1 & 0xffff0000u) >> 8u | (r2 & 0x000000ffu);
316    var z = to_float01(r12);
317    var r22 = r2 >> 8u;
318    var w = to_float01(r22);
319    return vec4<f32>(x, y, z, w);
320}
321
322fn rand_uniform_f(a: f32, b: f32) -> f32 {
323    return a + frand() * (b - a);
324}
325
326fn rand_uniform_vec2(a: vec2<f32>, b: vec2<f32>) -> vec2<f32> {
327    return a + frand2() * (b - a);
328}
329
330fn rand_uniform_vec3(a: vec3<f32>, b: vec3<f32>) -> vec3<f32> {
331    return a + frand3() * (b - a);
332}
333
334fn rand_uniform_vec4(a: vec4<f32>, b: vec4<f32>) -> vec4<f32> {
335    return a + frand4() * (b - a);
336}
337
338// Normal distribution computed using Box-Muller transform
339fn rand_normal_f(mean: f32, std_dev: f32) -> f32 {
340    var u = frand();
341    var v = frand();
342    var r = sqrt(-2.0 * log(u));
343    return mean + std_dev * r * cos(tau * v);
344}
345
346fn rand_normal_vec2(mean: vec2f, std_dev: vec2f) -> vec2f {
347    var u = frand();
348    var v = frand2();
349    var r = sqrt(-2.0 * log(u));
350    return mean + std_dev * r * cos(tau * v);
351}
352
353fn rand_normal_vec3(mean: vec3f, std_dev: vec3f) -> vec3f {
354    var u = frand();
355    var v = frand3();
356    var r = sqrt(-2.0 * log(u));
357    return mean + std_dev * r * cos(tau * v);
358}
359
360fn rand_normal_vec4(mean: vec4f, std_dev: vec4f) -> vec4f {
361    var u = frand();
362    var v = frand4();
363    var r = sqrt(-2.0 * log(u));
364    return mean + std_dev * r * cos(tau * v);
365}
366
367fn proj(u: vec3<f32>, v: vec3<f32>) -> vec3<f32> {
368    return dot(v, u) / dot(u,u) * u;
369}