zig/lib/std/compress/flate/BitWriter.zig

//! Bit writer for use in deflate (compression).
//!
//! Has internal bits buffer of 64 bits and internal bytes buffer of 248 bytes.
//! When we accumulate 48 bits 6 bytes are moved to the bytes buffer. When we
//! accumulate 240 bytes they are flushed to the underlying inner_writer.

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

// buffer_flush_size indicates the buffer size
// after which bytes are flushed to the writer.
// Should preferably be a multiple of 6, since
// we accumulate 6 bytes between writes to the buffer.
const buffer_flush_size = 240;

// buffer_size is the actual output byte buffer size.
// It must have additional headroom for a flush
// which can contain up to 8 bytes.
const buffer_size = buffer_flush_size + 8;

inner_writer: *std.io.BufferedWriter,

// Data waiting to be written is bytes[0 .. nbytes]
// and then the low nbits of bits.  Data is always written
// sequentially into the bytes array.
bits: u64 = 0,
nbits: u32 = 0, // number of bits
bytes: [buffer_size]u8 = undefined,
nbytes: u32 = 0, // number of bytes

const Self = @This();

pub fn init(bw: *std.io.BufferedWriter) Self {
    return .{ .inner_writer = bw };
}

pub fn setWriter(self: *Self, new_writer: *std.io.BufferedWriter) void {
    self.inner_writer = new_writer;
}

pub fn flush(self: *Self) anyerror!void {
    var n = self.nbytes;
    while (self.nbits != 0) {
        self.bytes[n] = @as(u8, @truncate(self.bits));
        self.bits >>= 8;
        if (self.nbits > 8) { // Avoid underflow
            self.nbits -= 8;
        } else {
            self.nbits = 0;
        }
        n += 1;
    }
    self.bits = 0;
    _ = try self.inner_writer.write(self.bytes[0..n]);
    self.nbytes = 0;
}

pub fn writeBits(self: *Self, b: u32, nb: u32) anyerror!void {
    self.bits |= @as(u64, @intCast(b)) << @as(u6, @intCast(self.nbits));
    self.nbits += nb;
    if (self.nbits < 48)
        return;

    var n = self.nbytes;
    std.mem.writeInt(u64, self.bytes[n..][0..8], self.bits, .little);
    n += 6;
    if (n >= buffer_flush_size) {
        _ = try self.inner_writer.write(self.bytes[0..n]);
        n = 0;
    }
    self.nbytes = n;
    self.bits >>= 48;
    self.nbits -= 48;
}

pub fn writeBytes(self: *Self, bytes: []const u8) anyerror!void {
    var n = self.nbytes;
    if (self.nbits & 7 != 0) {
        return error.UnfinishedBits;
    }
    while (self.nbits != 0) {
        self.bytes[n] = @as(u8, @truncate(self.bits));
        self.bits >>= 8;
        self.nbits -= 8;
        n += 1;
    }
    if (n != 0) {
        _ = try self.inner_writer.write(self.bytes[0..n]);
    }
    self.nbytes = 0;
    _ = try self.inner_writer.write(bytes);
}