vyre 0.4.0

const HEAP_CAPACITY: u32 = 128u;
const HEAP_OP_PUSH: u32 = 1u;
const HEAP_OP_POP_MAX: u32 = 2u;
const HEAP_OP_PEEK_MAX: u32 = 3u;
const HEAP_OK: u32 = 0u;
const HEAP_OVERFLOW: u32 = 1u;
const HEAP_UNDERFLOW: u32 = 2u;

struct HeapCommand {
    op: u32,
    lane: u32,
    value: u32,
    priority: u32,
}

struct HeapResult {
    status: u32,
    value: u32,
    priority: u32,
    len: u32,
}

@group(0) @binding(0) var<storage, read> heap_commands: array<HeapCommand>;
@group(0) @binding(1) var<storage, read_write> heap_results: array<HeapResult>;

var<workgroup> heap_values: array<u32, 128>;
var<workgroup> heap_priorities: array<u32, 128>;
var<workgroup> heap_len: atomic<u32>;

fn heap_higher(left: u32, right: u32) -> bool {
    let left_priority = heap_priorities[left];
    let right_priority = heap_priorities[right];
    let left_value = heap_values[left];
    let right_value = heap_values[right];
    return left_priority > right_priority ||
        (left_priority == right_priority && left_value > right_value);
}

fn heap_swap(left: u32, right: u32) {
    let value = heap_values[left];
    let priority = heap_priorities[left];
    heap_values[left] = heap_values[right];
    heap_priorities[left] = heap_priorities[right];
    heap_values[right] = value;
    heap_priorities[right] = priority;
}

fn heap_sift_up(start: u32) {
    var index = start;
    loop {
        if (index == 0u) {
            break;
        }
        let parent = (index - 1u) / 2u;
        if (!heap_higher(index, parent)) {
            break;
        }
        heap_swap(index, parent);
        index = parent;
    }
}

fn heap_sift_down(start: u32, count: u32) {
    var index = start;
    loop {
        let left = index * 2u + 1u;
        let right = left + 1u;
        var best = index;
        if (left < count && heap_higher(left, best)) {
            best = left;
        }
        if (right < count && heap_higher(right, best)) {
            best = right;
        }
        if (best == index) {
            break;
        }
        heap_swap(index, best);
        index = best;
    }
}

@compute @workgroup_size(64, 1, 1)
fn workgroup_queue_priority_kernel(@builtin(local_invocation_id) local_id: vec3<u32>) {
    let lane = local_id.x;
    if (lane == 0u) {
        atomicStore(&heap_len, 0u);
    }
    workgroupBarrier();

    var command_index = 0u;
    let command_count = arrayLength(&heap_commands);
    loop {
        if (command_index >= command_count) {
            break;
        }

        let command = heap_commands[command_index];
        if (command.lane == lane) {
            var status = HEAP_OK;
            var value = 0xffffffffu;
            var priority = 0xffffffffu;
            var len_after = atomicLoad(&heap_len);

            if (command.op == HEAP_OP_PUSH) {
                let slot = atomicAdd(&heap_len, 1u);
                if (slot >= HEAP_CAPACITY) {
                    _ = atomicSub(&heap_len, 1u);
                    status = HEAP_OVERFLOW;
                    len_after = HEAP_CAPACITY;
                } else {
                    heap_values[slot] = command.value;
                    heap_priorities[slot] = command.priority;
                    heap_sift_up(slot);
                    len_after = slot + 1u;
                }
            } else if (command.op == HEAP_OP_POP_MAX) {
                let count = atomicLoad(&heap_len);
                if (count == 0u) {
                    status = HEAP_UNDERFLOW;
                    len_after = 0u;
                } else {
                    value = heap_values[0];
                    priority = heap_priorities[0];
                    let last = count - 1u;
                    heap_values[0] = heap_values[last];
                    heap_priorities[0] = heap_priorities[last];
                    _ = atomicSub(&heap_len, 1u);
                    len_after = last;
                    if (last > 0u) {
                        heap_sift_down(0u, last);
                    }
                }
            } else if (command.op == HEAP_OP_PEEK_MAX) {
                let count = atomicLoad(&heap_len);
                if (count == 0u) {
                    status = HEAP_UNDERFLOW;
                    len_after = 0u;
                } else {
                    value = heap_values[0];
                    priority = heap_priorities[0];
                    len_after = count;
                }
            }

            heap_results[command_index] = HeapResult(status, value, priority, len_after);
        }

        workgroupBarrier();
        command_index = command_index + 1u;
    }
}