1struct KeyValuePair {
2    /// Sorting key.
3    key: u32,
4#ifdef HAS_DUAL_KEY
5    /// Secondary sorting key. Sorts value with the same primary key.
6    key2: u32,
7#endif
8    /// Value associated with the sort key(s), generally an index to some other data.
9    /// Copied as is and otherwise ignored by the sorting algorithm.
10    value: u32,
11}
12
13struct SortBuffer {
14    count: i32,
15    pairs: array<KeyValuePair>,
16}
17
18/// Check whether kv1 > kv2, comparing the key(s) of each pair.
19fn compare_greater(kv1: KeyValuePair, kv2: KeyValuePair) -> bool {
20    if (kv1.key > kv2.key) {
21        return true;
22    }
23#ifdef HAS_DUAL_KEY
24    if (kv1.key == kv2.key) {
25        return kv1.key2 > kv2.key2;
26    }
27#endif
28    return false;
29}
30
31@group(0) @binding(0) var<storage, read_write> sort_buffer : SortBuffer;
32
33/// Naive insertion sort. TODO: replace with something faster.
34@compute @workgroup_size(64)
35fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
36    // Naive single-threaded sort
37    if (global_invocation_id.x != 0) {
38        return;
39    }
40
41    // Insertion sort
42    let num_items = sort_buffer.count;
43    for (var i: i32 = 1; i < num_items; i += 1) {
44        var kv = sort_buffer.pairs[i];
45        var j = i;
46        while (j > 0 && compare_greater(sort_buffer.pairs[j - 1], kv)) {
47            sort_buffer.pairs[j] = sort_buffer.pairs[j - 1];
48            j -= 1;
49        }
50        sort_buffer.pairs[j] = kv;
51    }
52
53    // Clear for next frame
54    sort_buffer.count = 0;
55}