1#import bevy_render::view::View
2#import bevy_hanabi::vfx_common::{
3    IndirectBuffer, 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
13{{PROPERTIES}}
14
15struct ParticleBuffer {
16    particles: array<Particle>,
17}
18
19struct VertexOutput {
20    @builtin(position) position: vec4<f32>,
21    @location(0) color: vec4<f32>,
22#ifdef NEEDS_UV
23    @location(1) uv: vec2<f32>,
24#endif
25#ifdef NEEDS_NORMAL
26    @location(2) normal: vec3<f32>,
27#endif
28#ifdef NEEDS_PARTICLE_FRAGMENT
29    @location(3) slab_particle_index: u32,
30#endif
31}
32
33@group(0) @binding(0) var<uniform> view: View;
34@group(0) @binding(1) var<uniform> sim_params : SimParams;
35
36@group(1) @binding(0) var<storage, read> particle_buffer : ParticleBuffer;
37@group(1) @binding(1) var<storage, read> indirect_buffer : IndirectBuffer;
38
39// "spawner" group @2
40@group(2) @binding(0) var<storage, read> spawners : array<Spawner>;
41@group(2) @binding(1) var<storage, read> prefix_sum : array<u32>;
42@group(2) @binding(2) var<storage, read> batch_info : BatchInfo;
43{{PROPERTIES_BINDING}}
44
45{{MATERIAL_BINDINGS}}
46
47fn get_camera_position_effect_space() -> vec3<f32> {
48    let view_pos = view.world_from_view[3].xyz;
49#ifdef LOCAL_SPACE_SIMULATION
50    let inverse_transform = transpose(
51        mat3x3(
52            spawners[batch_info.base_effect + effect_location.effect_index].inverse_transform[0].xyz,
53            spawners[batch_info.base_effect + effect_location.effect_index].inverse_transform[1].xyz,
54            spawners[batch_info.base_effect + effect_location.effect_index].inverse_transform[2].xyz,
55        )
56    );
57    return inverse_transform * view_pos;
58#else
59    return view_pos;
60#endif
61}
62
63fn get_camera_rotation_effect_space() -> mat3x3<f32> {
64    let view_rot = mat3x3(view.world_from_view[0].xyz, view.world_from_view[1].xyz, view.world_from_view[2].xyz);
65#ifdef LOCAL_SPACE_SIMULATION
66    let inverse_transform = transpose(
67        mat3x3(
68            spawners[batch_info.base_effect + effect_location.effect_index].inverse_transform[0].xyz,
69            spawners[batch_info.base_effect + effect_location.effect_index].inverse_transform[1].xyz,
70            spawners[batch_info.base_effect + effect_location.effect_index].inverse_transform[2].xyz,
71        )
72    );
73    return inverse_transform * view_rot;
74#else
75    return view_rot;
76#endif
77}
78
79/// Unpack a compressed transform stored in transposed row-major form.
80fn unpack_compressed_transform(compressed_transform: mat3x4<f32>) -> mat4x4<f32> {
81    return transpose(
82        mat4x4(
83            compressed_transform[0],
84            compressed_transform[1],
85            compressed_transform[2],
86            vec4<f32>(0.0, 0.0, 0.0, 1.0)
87        )
88    );
89}
90
91// Unpacks a compressed transform and transposes is.
92fn unpack_compressed_transform_3x3_transpose(compressed_transform: mat3x4<f32>) -> mat3x3<f32> {
93    return mat3x3(
94        compressed_transform[0].xyz,
95        compressed_transform[1].xyz,
96        compressed_transform[2].xyz,
97    );
98}
99
100/// Transform a simulation space position into a world space position.
101///
102/// The simulation space depends on the effect's SimulationSpace value, and is either
103/// the effect space (SimulationSpace::Local) or the world space (SimulationSpace::Global).
104fn transform_position_simulation_to_world(sim_position: vec3<f32>) -> vec4<f32> {
105#ifdef LOCAL_SPACE_SIMULATION
106    let transform = unpack_compressed_transform(spawners[batch_info.base_effect + effect_location.effect_index].transform);
107    return transform * vec4<f32>(sim_position, 1.0);
108#else
109    return vec4<f32>(sim_position, 1.0);
110#endif
111}
112
113fn transform_normal_simulation_to_world(sim_normal: vec3<f32>) -> vec3<f32> {
114#ifdef LOCAL_SPACE_SIMULATION
115    // We use the inverse transpose transform to transform normals.
116    // The inverse transpose is the same as the transposed inverse, so we can
117    // safely use the inverse transform.
118    let transform = unpack_compressed_transform_3x3_transpose(spawners[batch_info.base_effect + effect_location.effect_index].inverse_transform);
119    return transform * sim_normal;
120#else
121    return sim_normal;
122#endif
123}
124
125/// Transform a simulation space position into a clip space position.
126///
127/// The simulation space depends on the effect's SimulationSpace value, and is either
128/// the effect space (SimulationSpace::Local) or the world space (SimulationSpace::Global).
129/// The clip space is the final [-1:1]^3 space output from the vertex shader, before
130/// perspective divide and viewport transform are applied.
131fn transform_position_simulation_to_clip(sim_position: vec3<f32>) -> vec4<f32> {
132    return view.clip_from_world * transform_position_simulation_to_world(sim_position);
133}
134
135fn inverse_transpose_mat3(m: mat3x3<f32>) -> mat3x3<f32> {
136    let tmp0 = cross(m[1], m[2]);
137    let tmp1 = cross(m[2], m[0]);
138    let tmp2 = cross(m[0], m[1]);
139    let inv_det = 1.0 / dot(m[2], tmp2);
140    return mat3x3<f32>(tmp0 * inv_det, tmp1 * inv_det, tmp2 * inv_det);
141}
142
143{{RENDER_EXTRA}}
144
145/// Location of an effect in a slab.
146struct EffectLocation {
147    /// Index of the effect in the global list of effects.
148    effect_index: u32,
149    /// Base particle index, that is index in the slab of the first particle for this instance.
150    base_particle: u32,
151    /// Index of this particle relative to its effect. Note that if there's an indirection
152    /// buffer then this is the linear index in [0:N[ of the particle to update, before the indirection.
153    update_index: u32,
154}
155
156/// Find the index of an effect from the index of a particle.
157///
158/// This uses a binary search on the slab_offset field of the spawners array, which
159/// represents a prefix sum of the particle count per effect (for previous effects;
160/// the value is actually the base particle so the first entry is always 0).
161///
162/// Requirements:
163/// - var<storage, read> batch_info : BatchInfo
164/// - var<storage, read> prefix_sum : array<u32>
165fn find_location_from_particle(slab_particle_index: u32) -> EffectLocation {
166    var lo = batch_info.prefix_sum_offset;
167    var hi = lo + batch_info.prefix_sum_count;
168    var num_iter = 0;  // avoid deadlocking the GPU by capping the iteration count
169    while (lo < hi) {
170        let mid = (hi + lo) >> 1u;
171        let base_particle = prefix_sum[mid];
172        if (slab_particle_index >= base_particle) {
173            lo = mid + 1u;
174        } else if (slab_particle_index < base_particle) {
175            hi = mid;
176        }
177        num_iter += 1;
178        if (num_iter >= 100) {
179            return EffectLocation(0xDEADBEEFu, 0xDEADBEEFu, 0xDEADBEEFu);
180        }
181    }
182    let base_particle = batch_info.base_particle + prefix_sum[lo - 1u];
183    let effect_index = lo - 1u - batch_info.prefix_sum_offset;
184    let update_index = slab_particle_index - base_particle;
185    return EffectLocation(effect_index, base_particle, update_index);
186}
187
188/// The resolved effect and particle location.
189///
190/// This is calculated at the start of the thread execution, and used after that
191/// in various functions.
192var<private> effect_location : EffectLocation;
193
194var<private> effect_metadata_index: u32;
195// var<private> properties_array_index: u32;
196
197@vertex
198fn vertex(
199    @builtin(instance_index) instance_index: u32,
200    @location(0) vertex_position: vec3<f32>,
201#ifdef NEEDS_UV
202    @location(1) vertex_uv: vec2<f32>,
203#endif
204#ifdef NEEDS_NORMAL
205    @location(2) vertex_normal: vec3<f32>,
206#endif
207    // @location(1) vertex_color: u32,
208    // @location(1) vertex_velocity: vec3<f32>,
209) -> VertexOutput {
210    // Global particle index into the slab, including those particles from other
211    // effect instances in the same batch, as well as possibly from other batches.
212    // This is rarely useful on its own.
213    let slab_particle_index = batch_info.base_particle + instance_index;
214
215    // Find the index of the effect this particle is part of.
216    effect_location = find_location_from_particle(slab_particle_index);
217    let spawner = &spawners[batch_info.base_effect + effect_location.effect_index];
218    effect_metadata_index = (*spawner).effect_metadata_index;
219    let base_particle = effect_location.base_particle;
220
221    // Fetch particle
222    let indirect_read_index = (*spawner).render_indirect_read_index;
223    let particle_index = indirect_buffer.rows[base_particle + instance_index].particle_index[indirect_read_index];
224    var particle = particle_buffer.particles[base_particle + particle_index];
225
226    var out: VertexOutput;
227
228#ifdef NEEDS_PARTICLE_FRAGMENT
229    out.slab_particle_index = base_particle + particle_index;
230#endif // NEEDS_PARTICLE_FRAGMENT
231
232#ifdef RIBBONS
233    // Discard first instance; we draw from second one, and link to previous one
234    if (instance_index == 0) {
235        out.position = vec4(0.0);
236        return out;
237    }
238
239    // Fetch previous particle
240    let prev_index = indirect_buffer.rows[base_particle + instance_index - 1u].particle_index[indirect_read_index];
241    let prev_particle = particle_buffer.particles[base_particle + prev_index];
242
243    // Discard this instance if previous one is from a different ribbon. Again,
244    // we draw from second one of each ribbon.
245    if (prev_particle.ribbon_id != particle.ribbon_id) {
246        out.position = vec4(0.0);
247        return out;
248    }
249#endif  // RIBBONS
250
251#ifdef NEEDS_UV
252    // Compute UVs
253    var uv = vertex_uv;
254#ifdef FLIPBOOK
255    let row_count = {{FLIPBOOK_ROW_COUNT}};
256    let ij = vec2<f32>(f32(particle.sprite_index % row_count), f32(particle.sprite_index / row_count));
257    uv = (ij + uv) * {{FLIPBOOK_SCALE}};
258#endif
259    out.uv = uv;
260#endif  // NEEDS_UV
261
262{{INPUTS}}
263
264{{VERTEX_MODIFIERS}}
265
266#ifdef RIBBONS
267    var delta = particle.position - prev_particle.position;
268
269    axis_x = normalize(delta);
270    axis_y = normalize(cross(axis_x, axis_z));
271    axis_z = cross(axis_x, axis_y);
272
273    position = mix(particle.position, prev_particle.position, 0.5);
274    size = vec3(length(delta), size.y, 1.0);
275#endif  // RIBBONS
276
277    // Expand particle mesh vertex based on particle position ("origin"), and local
278    // orientation and size of the particle mesh.
279    let vpos = vertex_position * size;
280    let sim_position = position + axis_x * vpos.x + axis_y * vpos.y + axis_z * vpos.z;
281    out.position = transform_position_simulation_to_clip(sim_position);
282
283    out.color = color;
284
285#ifdef NEEDS_NORMAL
286    let normal = inverse_transpose_mat3(mat3x3(axis_x, axis_y, axis_z)) * vertex_normal;
287    out.normal = transform_normal_simulation_to_world(normal);
288#endif  // NEEDS_NORMAL
289
290    return out;
291}
292
293@fragment
294fn fragment(in: VertexOutput) -> @location(0) vec4<f32> {
295    // Read fragment inputs
296#ifdef USE_ALPHA_MASK
297    var alpha_cutoff: f32 = {{ALPHA_CUTOFF}};
298#endif
299    var color = in.color;
300#ifdef NEEDS_UV
301    var uv = in.uv;
302#endif
303#ifdef NEEDS_NORMAL
304    var normal = in.normal;
305#endif
306#ifdef NEEDS_PARTICLE_FRAGMENT
307    var particle = particle_buffer.particles[in.slab_particle_index];
308#endif // NEEDS_PARTICLE_FRAGMENT
309
310{{FRAGMENT_MODIFIERS}}
311
312#ifdef USE_ALPHA_MASK
313    if color.a >= alpha_cutoff {
314        color.a = 1.0;
315    } else {
316        discard;
317    }
318#endif
319
320    return color;
321}