vyre 0.4.0

const HASHMAP_CAPACITY: u32 = 1024u;
const HASHMAP_EMPTY_KEY: u32 = 0xffffffffu;
const HASHMAP_EMPTY_VALUE: u32 = 0xffffffffu;
const HASHMAP_OP_INSERT: u32 = 1u;
const HASHMAP_OP_LOOKUP: u32 = 2u;
const HASHMAP_OP_REMOVE: u32 = 3u;
const HASHMAP_OK: u32 = 0u;
const HASHMAP_ERR_OVERFLOW: u32 = 1u;
const HASHMAP_ERR_NOT_FOUND: u32 = 2u;
const HASHMAP_REPLACED: u32 = 3u;
const HASHMAP_ERR_RESERVED_KEY: u32 = 4u;
const HASHMAP_ERR_INVALID_CAPACITY: u32 = 5u;

struct HashmapCommand {
    op: u32,
    lane: u32,
    key: u32,
    value: u32,
}

struct HashmapResult {
    status: u32,
    value: u32,
    len: u32,
    reserved: u32,
}

@group(0) @binding(0) var<storage, read> hashmap_commands: array<HashmapCommand>;
@group(0) @binding(1) var<storage, read_write> hashmap_results: array<HashmapResult>;

var<workgroup> hashmap_keys: array<u32, 1024>;
var<workgroup> hashmap_values: array<u32, 1024>;
var<workgroup> hashmap_len: u32;

fn hash_u32(key: u32) -> u32 {
    var x = key + 0x9e3779b1u;
    x = (x ^ (x >> 16u)) * 0x7feb352du;
    x = (x ^ (x >> 15u)) * 0x846ca68bu;
    return x ^ (x >> 16u);
}

fn hashmap_probe_start(key: u32) -> u32 {
    return hash_u32(key) & (HASHMAP_CAPACITY - 1u);
}

fn hashmap_find_slot(key: u32) -> u32 {
    let start = hashmap_probe_start(key);
    var probe = 0u;
    loop {
        if (probe >= HASHMAP_CAPACITY) {
            return HASHMAP_EMPTY_KEY;
        }
        let idx = (start + probe) & (HASHMAP_CAPACITY - 1u);
        let slot_key = hashmap_keys[idx];
        if (slot_key == HASHMAP_EMPTY_KEY || slot_key == key) {
            return idx;
        }
        probe = probe + 1u;
    }
}

fn hashmap_insert(key: u32, value: u32) -> u32 {
    let idx = hashmap_find_slot(key);
    if (idx == HASHMAP_EMPTY_KEY) {
        return HASHMAP_ERR_OVERFLOW;
    }
    if (hashmap_keys[idx] == key) {
        hashmap_values[idx] = value;
        return HASHMAP_REPLACED;
    }
    hashmap_keys[idx] = key;
    hashmap_values[idx] = value;
    hashmap_len = hashmap_len + 1u;
    return HASHMAP_OK;
}

fn hashmap_lookup(key: u32) -> HashmapResult {
    let start = hashmap_probe_start(key);
    var probe = 0u;
    loop {
        if (probe >= HASHMAP_CAPACITY) {
            return HashmapResult(HASHMAP_ERR_NOT_FOUND, HASHMAP_EMPTY_VALUE, hashmap_len, 0u);
        }
        let idx = (start + probe) & (HASHMAP_CAPACITY - 1u);
        let slot_key = hashmap_keys[idx];
        if (slot_key == HASHMAP_EMPTY_KEY) {
            return HashmapResult(HASHMAP_ERR_NOT_FOUND, HASHMAP_EMPTY_VALUE, hashmap_len, 0u);
        }
        if (slot_key == key) {
            return HashmapResult(HASHMAP_OK, hashmap_values[idx], hashmap_len, 0u);
        }
        probe = probe + 1u;
    }
}

fn hashmap_remove(key: u32) -> HashmapResult {
    let start = hashmap_probe_start(key);
    var probe = 0u;
    var hit = HASHMAP_EMPTY_KEY;
    var removed_value = HASHMAP_EMPTY_VALUE;
    loop {
        if (probe >= HASHMAP_CAPACITY) {
            return HashmapResult(HASHMAP_ERR_NOT_FOUND, HASHMAP_EMPTY_VALUE, hashmap_len, 0u);
        }
        let idx = (start + probe) & (HASHMAP_CAPACITY - 1u);
        let slot_key = hashmap_keys[idx];
        if (slot_key == HASHMAP_EMPTY_KEY) {
            return HashmapResult(HASHMAP_ERR_NOT_FOUND, HASHMAP_EMPTY_VALUE, hashmap_len, 0u);
        }
        if (slot_key == key) {
            hit = idx;
            removed_value = hashmap_values[idx];
            break;
        }
        probe = probe + 1u;
    }

    hashmap_keys[hit] = HASHMAP_EMPTY_KEY;
    hashmap_values[hit] = HASHMAP_EMPTY_VALUE;
    hashmap_len = hashmap_len - 1u;

    // Backward-shift deletion: reinsert displaced entries so probe chains
    // remain searchable after removing a middle slot.
    var cursor = (hit + 1u) & (HASHMAP_CAPACITY - 1u);
    loop {
        if (hashmap_keys[cursor] == HASHMAP_EMPTY_KEY) {
            break;
        }
        let displaced_key = hashmap_keys[cursor];
        let displaced_value = hashmap_values[cursor];
        hashmap_keys[cursor] = HASHMAP_EMPTY_KEY;
        hashmap_values[cursor] = HASHMAP_EMPTY_VALUE;
        hashmap_len = hashmap_len - 1u;
        _ = hashmap_insert(displaced_key, displaced_value);
        cursor = (cursor + 1u) & (HASHMAP_CAPACITY - 1u);
    }
    return HashmapResult(HASHMAP_OK, removed_value, hashmap_len, 0u);
}

@compute @workgroup_size(64, 1, 1)
fn workgroup_hashmap_kernel(@builtin(local_invocation_id) local_id: vec3<u32>) {
    let lane = local_id.x;
    var init_index = lane;
    loop {
        if (init_index >= HASHMAP_CAPACITY) {
            break;
        }
        hashmap_keys[init_index] = HASHMAP_EMPTY_KEY;
        hashmap_values[init_index] = HASHMAP_EMPTY_VALUE;
        init_index = init_index + 64u;
    }
    if (lane == 0u) {
        hashmap_len = 0u;
    }
    workgroupBarrier();

    var command_index = 0u;
    let command_count = arrayLength(&hashmap_commands);
    loop {
        if (command_index >= command_count) {
            break;
        }
        let command = hashmap_commands[command_index];
        if (command.lane == lane) {
            var result = HashmapResult(HASHMAP_OK, HASHMAP_EMPTY_VALUE, hashmap_len, 0u);
            if (command.op == HASHMAP_OP_INSERT) {
                let status = hashmap_insert(command.key, command.value);
                result = HashmapResult(status, HASHMAP_EMPTY_VALUE, hashmap_len, 0u);
            } else if (command.op == HASHMAP_OP_LOOKUP) {
                result = hashmap_lookup(command.key);
            } else if (command.op == HASHMAP_OP_REMOVE) {
                result = hashmap_remove(command.key);
            }
            hashmap_results[command_index] = result;
        }
        workgroupBarrier();
        command_index = command_index + 1u;
    }
}