Packed array of enum

Help
1

I’m trying to do more DoD, but struggling with some ergonomic issues.

I have a const Spec = enum(u4) {...} and I want to store 8 of it. What is the best option ?
[8]Spec seems to defeat the purpose of packing stuff in a u4 in the first place,
@Vector(Spec, 8) doesn’t compile because Spec isn’t an int.

So I wrote this ugly code:

pub const Array = packed struct {
    _0: Spec,
    _1: Spec,
    _2: Spec,
    _3: Spec,
    _4: Spec,
    _5: Spec,
    _6: Spec,
    _7: Spec,

    pub const unknown: Array = splat(.unknown);
    const MAX_RANK = 8;
    const Vec = @Vector(MAX_RANK, u4);

    pub fn init(specs: []const Spec) Array {
        var res: Array = unknown;
        for (0.., specs) |i, spec| {
            res = res.set(i, spec);
        }
        return res;
    }

    pub fn replicated(rank_: usize) Array {
        std.debug.assert(rank_ <= MAX_RANK);
        const full_replicated: Vec = @splat(@intFromEnum(Spec.replicated));
        const mask = std.simd.iota(u4, 8) < @as(Vec, @splat(@truncate(rank_)));
        return @bitCast(@select(u4, mask, full_replicated, @as(Vec, @bitCast(unknown))));
    }

    pub fn splat(spec: Spec) Array {
        const vec: Vec = @splat(@intFromEnum(spec));
        return @bitCast(vec);
    }

    pub fn get(array: Array, ax: usize) Spec {
        std.debug.assert(ax < MAX_RANK);
        const pack: u32 = @bitCast(array);
        const shift: u5 = @intCast(4 * ax);
        return @enumFromInt(@as(u4, @truncate(pack >> shift)));
    }

    pub fn set(array: Array, ax: usize, spec: Spec) Array {
        std.debug.assert(ax < MAX_RANK);
        const pack: u32 = @bitCast(array);
        const shift: u5 = @intCast(4 * ax);
        const mask = @as(u32, 0xf) << shift;
        return @bitCast((pack & ~mask) | (@as(u32, @intFromEnum(spec)) << shift));
    }

and I’m not even including the “insert” method.

Is there smth better ? what would you recommend here ?

2

I use this to represent arrays of a bitpacked type (which I use A LOT for mmio types)

/// Represents a bitpacked array of `n` elements of `T`.
/// Stored in native endian.
///
/// A `BitpackedArray` is stored from LSb (0) to MSb (n - 1).
pub fn BitpackedArray(comptime T: type, comptime n: usize) type {
    const total_bit_size = @bitSizeOf(T) * n;
    const ArrayInt = @Int(.unsigned, total_bit_size);
    const ElementInt = @Int(.unsigned, @bitSizeOf(T));

    return packed struct(ArrayInt) {
        pub const Int = ArrayInt;
        pub const elem_mask: ArrayInt = std.math.maxInt(ElementInt);

        raw: ArrayInt,

        pub fn init(value: [n]T) Self {
            // NOTE: Cannot be `undefined`, any `undefined` bits make the entire value `undefined`.
            var bt: Self = std.mem.zeroes(Self);
            inline for (0..n) |i| bt.set(i, value[i]);
            return bt;
        }

        pub fn splat(value: T) Self {
            // NOTE: Cannot be `undefined`, any `undefined` bits make the entire value `undefined`.
            var bt: Self = std.mem.zeroes(Self);
            inline for (0..n) |i| bt.set(i, value);
            return bt;
        }

        pub fn slice(bt: Self, index: usize, comptime len: usize) BitpackedArray(T, len) {
            std.debug.assert(index + len <= n);

            const NewBitpacked = BitpackedArray(T, len);
            const bt_int: Int = @bitCast(bt);
            const new_bt_int: NewBitpacked.Int = @truncate(bt_int >> (index * @bitSizeOf(T)));
            return @bitCast(new_bt_int);
        }

        pub fn get(bt: Self, index: usize) T {
            std.debug.assert(index < n);
            // TODO: @fromBackingInt 0.17.0
            const value: ElementInt = @truncate((@as(ArrayInt, @bitCast(bt)) >> @intCast(index * @bitSizeOf(T))) & elem_mask);

            return switch (@typeInfo(T)) {
                .@"enum" => @enumFromInt(value),
                else => @bitCast(value),
            };
        }

        pub fn copyWith(bt: Self, index: usize, value: T) Self {
            std.debug.assert(index < n);
            // TODO: @fromBackingInt 0.17.0
            const int: ElementInt = switch (@typeInfo(T)) {
                .@"enum" => @intFromEnum(value),
                else => @bitCast(value),
            };

            return @bitCast((@as(ArrayInt, @bitCast(bt)) & ~(elem_mask << @intCast(index * @bitSizeOf(T)))) | (@as(ArrayInt, int) << @intCast(index * @bitSizeOf(T))));
        }

        pub fn set(bt: anytype, index: usize, value: T) void {
            comptime std.debug.assert(@typeInfo(@TypeOf(bt)) == .pointer);
            comptime std.debug.assert(!@typeInfo(@TypeOf(bt)).pointer.is_const);
            comptime std.debug.assert(@typeInfo(@TypeOf(bt)).pointer.child == Self);
            std.debug.assert(index < n);

            bt.* = bt.copyWith(index, value);
        }

        pub fn format(bt: Self, w: *std.Io.Writer) std.Io.Writer.Error!void {
            try w.writeAll("{ ");
            for (0..n) |i| {
                try w.print(if (std.meta.hasFn(T, "format")) "{f}" else "{any}", .{bt.get(i)});
                if (i != n - 1) try w.writeAll(", ");
            }
            try w.writeAll(" }");
        }

        const Self = @This();
    };
}

test BitpackedArray {
    const Thing = enum(u1) { foo, bar };
    const ThingArray = BitpackedArray(Thing, 4);

    var bt: ThingArray = .splat(.foo);

    bt.set(3, .bar);

    try testing.expect(bt.get(3) == .bar);

    bt.set(3, .foo);

    try testing.expect(bt.get(3) == .foo);

    bt.set(0, .bar);

    try testing.expect(bt.get(0) == .bar);
}
1 Like
3

somehow I feel this should be part of the language or at least the std.
Thanks for sharing

1 Like