Comptime enum-based 1D-4D vector swizzling

swizzle.zig

fn swizzleEnum(comptime N: usize, comptime T: type, a: @Vector(N, T), comptime id: Swizzle) @Vector(swizzleLength(id), T) {
    comptime var code = @intFromEnum(id);
    comptime var i = 0;
    comptime var offset = 0;
    comptime var count = 4;
    inline while (code >= offset + count) {
        offset += count;
        count <<= 2;
        i += 1;
    }
    var res: [i + 1]T = undefined;
    code -= offset;
    inline while (i >= 0) : (i -= 1) {
        const idx = code & 3;
        if (idx >= N) @compileError("illegal vector swizzle: ." ++ @tagName(id));
        res[i] = a[idx];
        code >>= 2;
    }
    return res;
}

fn swizzleLength(id: Swizzle) usize {
    const i: usize = @intFromEnum(id);
    return if (i < 4) 1 else if (i < 20) 2 else if (i < 84) 3 else 4;
}

test "swizzleEnum" {
    try std.testing.expectEqual(
        [_]f32{ 2, 2, 1, 1 },
        swizzleEnum(2, f32, [_]f32{ 1, 2 }, .yyxx),
    );
}

const Swizzle = enum(u16) {
    x, // 0
    y,
    z,
    w,

    xx, // 4
    xy,
    xz,
    xw,
    yx,
    yy,
    yz,
    yw,
    zx,
    zy,
    zz,
    zw,
    wx,
    wy,
    wz,
    ww,

    xxx, // 20
    xxy,
    xxz,
    xxw,
    xyx,
    xyy,
    xyz,
    xyw,
    xzx,
    xzy,
    xzz,
    xzw,
    xwx,
    xwy,
    xwz,
    xww,
    yxx,
    yxy,
    yxz,
    yxw,
    yyx,
    yyy,
    yyz,
    yyw,
    yzx,
    yzy,
    yzz,
    yzw,
    ywx,
    ywy,
    ywz,
    yww,
    zxx,
    zxy,
    zxz,
    zxw,
    zyx,
    zyy,
    zyz,
    zyw,
    zzx,
    zzy,
    zzz,
    zzw,
    zwx,
    zwy,
    zwz,
    zww,
    wxx,
    wxy,
    wxz,
    wxw,
    wyx,
    wyy,
    wyz,
    wyw,
    wzx,
    wzy,
    wzz,
    wzw,
    wwx,
    wwy,
    wwz,
    www,

    xxxx, // 84
    xxxy,
    xxxz,
    xxxw,
    xxyx,
    xxyy,
    xxyz,
    xxyw,
    xxzx,
    xxzy,
    xxzz,
    xxzw,
    xxwx,
    xxwy,
    xxwz,
    xxww,
    xyxx,
    xyxy,
    xyxz,
    xyxw,
    xyyx,
    xyyy,
    xyyz,
    xyyw,
    xyzx,
    xyzy,
    xyzz,
    xyzw,
    xywx,
    xywy,
    xywz,
    xyww,
    xzxx,
    xzxy,
    xzxz,
    xzxw,
    xzyx,
    xzyy,
    xzyz,
    xzyw,
    xzzx,
    xzzy,
    xzzz,
    xzzw,
    xzwx,
    xzwy,
    xzwz,
    xzww,
    xwxx,
    xwxy,
    xwxz,
    xwxw,
    xwyx,
    xwyy,
    xwyz,
    xwyw,
    xwzx,
    xwzy,
    xwzz,
    xwzw,
    xwwx,
    xwwy,
    xwwz,
    xwww,

    yxxx,
    yxxy,
    yxxz,
    yxxw,
    yxyx,
    yxyy,
    yxyz,
    yxyw,
    yxzx,
    yxzy,
    yxzz,
    yxzw,
    yxwx,
    yxwy,
    yxwz,
    yxww,
    yyxx,
    yyxy,
    yyxz,
    yyxw,
    yyyx,
    yyyy,
    yyyz,
    yyyw,
    yyzx,
    yyzy,
    yyzz,
    yyzw,
    yywx,
    yywy,
    yywz,
    yyww,
    yzxx,
    yzxy,
    yzxz,
    yzxw,
    yzyx,
    yzyy,
    yzyz,
    yzyw,
    yzzx,
    yzzy,
    yzzz,
    yzzw,
    yzwx,
    yzwy,
    yzwz,
    yzww,
    ywxx,
    ywxy,
    ywxz,
    ywxw,
    ywyx,
    ywyy,
    ywyz,
    ywyw,
    ywzx,
    ywzy,
    ywzz,
    ywzw,
    ywwx,
    ywwy,
    ywwz,
    ywww,

    zxxx,
    zxxy,
    zxxz,
    zxxw,
    zxyx,
    zxyy,
    zxyz,
    zxyw,
    zxzx,
    zxzy,
    zxzz,
    zxzw,
    zxwx,
    zxwy,
    zxwz,
    zxww,
    zyxx,
    zyxy,
    zyxz,
    zyxw,
    zyyx,
    zyyy,
    zyyz,
    zyyw,
    zyzx,
    zyzy,
    zyzz,
    zyzw,
    zywx,
    zywy,
    zywz,
    zyww,
    zzxx,
    zzxy,
    zzxz,
    zzxw,
    zzyx,
    zzyy,
    zzyz,
    zzyw,
    zzzx,
    zzzy,
    zzzz,
    zzzw,
    zzwx,
    zzwy,
    zzwz,
    zzww,
    zwxx,
    zwxy,
    zwxz,
    zwxw,
    zwyx,
    zwyy,
    zwyz,
    zwyw,
    zwzx,
    zwzy,
    zwzz,
    zwzw,
    zwwx,
    zwwy,
    zwwz,
    zwww,

    wxxx,
    wxxy,
    wxxz,
    wxxw,
    wxyx,
    wxyy,
    wxyz,
    wxyw,
    wxzx,
    wxzy,
    wxzz,
    wxzw,
    wxwx,
    wxwy,
    wxwz,
    wxww,
    wyxx,
    wyxy,
    wyxz,
    wyxw,
    wyyx,
    wyyy,
    wyyz,
    wyyw,
    wyzx,
    wyzy,
    wyzz,
    wyzw,
    wywx,
    wywy,
    wywz,
    wyww,
    wzxx,
    wzxy,
    wzxz,
    wzxw,
    wzyx,
    wzyy,
    wzyz,
    wzyw,
    wzzx,
    wzzy,
    wzzz,
    wzzw,
    wzwx,
    wzwy,
    wzwz,
    wzww,
    wwxx,
    wwxy,
    wwxz,
    wwxw,
    wwyx,
    wwyy,
    wwyz,
    wwyw,
    wwzx,
    wwzy,
    wwzz,
    wwzw,
    wwwx,
    wwwy,
    wwwz,
    wwww,
};