[[vk::binding(0, 0)]]
RWStructuredBuffer<uint32_t> count_buffer;

[[vk::binding(0, 1)]]
RWStructuredBuffer<uint32_t> reduced_buffer;
[[vk::binding(1, 1)]]
RWStructuredBuffer<uint32_t> sum_buffer;
[[vk::binding(2, 1)]]
RWStructuredBuffer<uint32_t> compaction_buffer;

groupshared uint32_t local_data[256 * 2];
static uint32_t      THREAD_WIDTH = 256;
static uint32_t      DATA_WIDTH = THREAD_WIDTH * 2;

[numthreads(256, 1, 1)]
[shader("compute")]
void block_sum(
    uint32_t3 workgroup_id: SV_GroupID,
    uint32_t3 local_thread_id: SV_GroupThreadID,
    uint32_t3 global_thread_id: SV_DispatchThreadID)
{
    // Perform sum in current block

    // Copy local_datainto LDS with predicate
    let thread_index = global_thread_id.x;
    let local_thread_index = local_thread_id.x;
    let total = count_buffer.getCount();

    if (thread_index * 2 < total)
    {
        local_data[local_thread_index * 2] = select(count_buffer[thread_index * 2] != 0, 1, 0);
    }
    else
    {
        local_data[local_thread_index * 2] = 0;
    }

    if (thread_index * 2 + 1 < total)
    {
        local_data[local_thread_index * 2 + 1] = select(count_buffer[thread_index * 2 + 1] != 0, 1, 0);
    }
    else
    {
        local_data[local_thread_index * 2 + 1] = 0;
    }

    GroupMemoryBarrierWithGroupSync();

    var width : uint32_t = 2;
    while (width <= DATA_WIDTH)
    {
        let dest_index = width * (thread_index + 1) - 1;
        let get_index = dest_index - (width / 2);
        // println!("{}, {}", get_index, dest_index);
        if (dest_index < DATA_WIDTH)
        {
            local_data[dest_index] += local_data[get_index];
        }
        width *= 2;
        GroupMemoryBarrierWithGroupSync();
    }

    local_data[DATA_WIDTH - 1] = 0;
    while (width >= 2)
    {
        let dest_index = width * (thread_index + 1) - 1;
        let get_index = dest_index - (width / 2);
        // println!("{}, {}", get_index, dest_index);
        if (dest_index < DATA_WIDTH)
        {
            let self_data = local_data[dest_index];
            local_data[dest_index] += local_data[get_index];
            local_data[get_index] = self_data;
        }
        width /= 2;
        GroupMemoryBarrierWithGroupSync();
    }

    // Block now contains running local sum
    // Dump back to sum buffer
    sum_buffer[2 * thread_index] = local_data[2 * local_thread_index];
    sum_buffer[2 * thread_index + 1] = local_data[2 * local_thread_index + 1];

    // Write to reduced buffer
    reduced_buffer[workgroup_id.x] = local_data[DATA_WIDTH - 1];
}

[numthreads(1, 1, 1)]
[shader("compute")]
void linear_reduced_sum(
    uint32_t3 workgroup_id: SV_GroupID,
    uint32_t3 local_thread_id: SV_GroupThreadID,
    uint32_t3 global_thread_id: SV_DispatchThreadID)
{
    let size = reduced_buffer.getCount();

    // Perform exclusive sum
    var running_sum : uint32_t = 0;
    for (uint32_t i = 0; i < size; i++)
    {
        let value = reduced_buffer[i];
        reduced_buffer[i] = running_sum;
        running_sum += value;
    }
}

[numthreads(256, 1, 1)]
[shader("compute")]
void uniform_add(
    uint32_t3 workgroup_id: SV_GroupID,
    uint32_t3 local_thread_id: SV_GroupThreadID,
    uint32_t3 global_thread_id: SV_DispatchThreadID)
{
    let thread_index = global_thread_id.x;
    let local_thread_index = local_thread_id.x;

    // Gather
    let reduced_value = reduced_buffer[global_thread_id.x];
    // Apply
    sum_buffer[thread_index * 2] += reduced_value;
    sum_buffer[thread_index * 2 + 1] += reduced_value;
}