Zig Vector Ternary If

vector_ternary_if.zig

const assert = @import("std").debug.assert;

fn make_mask_3_f32(cond: [3]bool) [3]f32 {
    var mask: [3]f32 = undefined;
    const ar_cond: [3]bool = cond;

    comptime var i = 0;
    inline while (i < 3) : (i += 1) {
        mask[i] = if (ar_cond[i]) 1.0 else 0.0;
    }
    return mask;
}

fn ternary_if_3_f32(cond: [3]bool, a: @Vector(3, f32), b: @Vector(3, f32)) @Vector(3, f32) {
    const mask = make_mask_3_f32(cond); // bool to 1.0 or 0.0
    const vmask: @Vector(3, f32) = mask; // array to vector
    const ones: @Vector(3, f32) = [3]f32{ 1.0, 1.0, 1.0 }; // used to negate mask

    return vmask * a + (ones - vmask) * b;
}

test "cond to mask" {
    const cond = [3]bool{ true, false, true };
    const mask = make_mask_3_f32(cond);

    assert(mask[0] == 1.0);
    assert(mask[1] == 0.0);
    assert(mask[2] == 1.0);
}

test "ternary if" {
    const cond = [3]bool{ true, false, true };
    const a: @Vector(3, f32) = [3]f32{ 1.0, 2.0, 3.0 };
    const b: @Vector(3, f32) = [3]f32{ 10.0, 9.0, 8.0 };

    const result = ternary_if_3_f32(cond, a, b);

    assert(result[0] == 1.0);
    assert(result[1] == 9.0);
    assert(result[2] == 3.0);
}