1#import bevy_hanabi::vfx_common::{BatchInfo, DispatchIndirectArgs}
2
3@group(0) @binding(0) var<storage, read_write> batch_infos : array<BatchInfo>;
4@group(0) @binding(1) var<storage, read_write> prefix_sum : array<u32>;
5@group(0) @binding(2) var<storage, read_write> dispatch_indirect_buffer : array<DispatchIndirectArgs>;
6
7/// Compute the prefix sum of each effect batch.
8///
9/// This is invoked once per effect batch each frame, before the update pass, to recalculate
10/// the prefix sum of the alive particle counts per instance in that batch. Before this pass,
11/// the prefix sum buffer contins the number of alive particles for each instance. After it,
12/// it contains a prefix sum of that number.
13@compute @workgroup_size(64)
14fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
15    // Get the batch to update
16    let batch_index = global_invocation_id.x;
17    let batch_count = arrayLength(&batch_infos);
18    if (batch_index >= batch_count) {
19        return;
20    }
21
22    // Find the slice of prefix sum to update inside the global shared buffer
23    let offset = batch_infos[batch_index].prefix_sum_offset;
24    let count = batch_infos[batch_index].prefix_sum_count;
25
26    // Calculate the prefix sum
27    let end = offset + count;
28    var sum = 0u;
29    for (var i = offset; i < end; i += 1u) {
30        let count = prefix_sum[i];
31        // Exclusive: store current sum before incrementing
32        prefix_sum[i] = sum;
33        sum = sum + count;
34    }
35
36    // Store the sum back into the batch
37    batch_infos[batch_index].total_update_count = sum;
38
39    // Update the dispatch indirect count
40    dispatch_indirect_buffer[batch_index].x = (sum + 63u) >> 6u;
41    dispatch_indirect_buffer[batch_index].y = 1u;
42    dispatch_indirect_buffer[batch_index].z = 1u;
43}