2025-05-17 # Bit Stream

In the quest to make a good context-based adaptive arithmetic coder, I required an abstraction that could read bits from a byte source, and an abstraction that could write bits to a byte sink. Turns out, the code for this is pretty clean! The same two problems can be solved with only one structure.

The abstraction is pretty convenient - not only can I treat a byte source/sink pair as a unified bit stream, but the code generation for this is nearly optimal as well!

/// bit reader/writer wrapping a byte reader/writer
pub const BitStream = struct {
    buffer: u8,
    bits: u3,

    pub const init: @This() = .{ .bits = 0, .buffer = undefined };

    // read a single bit
    pub fn read(self: *@This(), reader: anytype) !u1 {
        if (self.bits == 0) {
            @branchHint(.unlikely);
            self.buffer = try reader.readByte();
        }
        self.bits -%= 1;
        return @truncate(self.buffer >> self.bits);
    }

    // write a single bit
    pub fn write(self: *@This(), writer: anytype, bit: u1) !void {
        const select: u8 = @as(u8, 1) << (7 - self.bits);
        self.buffer = (self.buffer & ~select) | (select * bit);
        if (self.bits == 7) {
            @branchHint(.unlikely);
            try writer.writeByte(self.buffer);
        }
        self.bits +%= 1;
    }

    // write all buffered bits (and maybe some undefined ones)
    pub fn flush(self: *@This(), writer: anytype) !void {
        if (self.bits != 0) {
            try writer.writeByte(self.buffer);
        }
    }
};

It gives me a good feeling when a simple idea just *works*. My implementation doesn't need to handle endianness, nor does it ever expect to read or write more than one bit at a time. Due to this, it's much cleaner than the current zig standard library bit reader and bit writer.