Bit-wise reader for Zig

bitreader.zig

const std = @import("std");

pub fn BitReader(comptime UnderlyingReader: type) type {
    return struct {
        const Self = @This();

        underlying_reader: UnderlyingReader,

        current_byte: u8 = 0,
        offset_into_current_byte: u8 = 8,

        pub fn readInt(self: *Self, comptime T: type) !T {
            const bits: u16 = @typeInfo(T).Int.bits;
            const UnsignedT = @Type(.{ .Int = .{ .signedness = .unsigned, .bits = bits }});

            var uvalue: UnsignedT = 0;

            var bit_offset: u16 = 0;

            while (bit_offset < bits) {
                if (self.offset_into_current_byte >= 8) {
                    self.current_byte = try self.underlying_reader.readByte();
                    self.offset_into_current_byte = 0;
                }

                const bits_satisfied_by_current_byte = std.math.min(bits - bit_offset, 8 - self.offset_into_current_byte);

                const mask = (@as(u8, 1) <<| bits_satisfied_by_current_byte) - 1;
                uvalue |=  @intCast(UnsignedT, (self.current_byte >> @intCast(u3, self.offset_into_current_byte)) & mask) <<| bit_offset;

                self.offset_into_current_byte += bits_satisfied_by_current_byte;
                bit_offset += bits_satisfied_by_current_byte;
            }

            return @bitCast(T, uvalue);
        }
    };
}

test "BitReader" {
    // 0001 0010 0011 0100 0101 0110 0111 1000
    const bits: u32 = 0x12345678;
    const bytes = std.mem.asBytes(&bits);

    var stream = std.io.fixedBufferStream(bytes);
    const reader = stream.reader();

    var bit_reader = BitReader(@TypeOf(reader)){ .underlying_reader = reader };

    // 0001 0010 0011 0100 0101 0110 0111 1[000]
    const a: u3 = try bit_reader.readInt(u3);
    try std.testing.expectEqual(@as(u3, 0b000), a);

    // 0001 0010 0011 0100 0101 0110 011[1 1]000
    const b: u2 = try bit_reader.readInt(u2);
    try std.testing.expectEqual(@as(u2, 0b11), b);

    // 0001 0010 0011 0100 01[01 0110 011]1 1000
    const c: u9 = try bit_reader.readInt(u9);
    try std.testing.expectEqual(@as(u9, 0b010110011), c);

    // 0001 001[0 0011 0100 01]01 0110 0111 1000
    const d: u11 = try bit_reader.readInt(u11);
    try std.testing.expectEqual(@as(u11, 0b00011010001), d);

    // [0001 001]0 0011 0100 0101 0110 0111 1000
    const e: u7 = try bit_reader.readInt(u7);
    try std.testing.expectEqual(@as(u7, 0b0001001), e);

    try std.testing.expectError(error.EndOfStream, bit_reader.readInt(u1));
}