| ... | @@ -1,179 +0,0 @@ |
| 1 | const std = @import("../std.zig"); |
| 2 | |
| 3 | //General note on endianess: |
| 4 | //Big endian is packed starting in the most significant part of the byte and subsequent |
| 5 | // bytes contain less significant bits. Thus we write out bits from the high end |
| 6 | // of our input first. |
| 7 | //Little endian is packed starting in the least significant part of the byte and |
| 8 | // subsequent bytes contain more significant bits. Thus we write out bits from |
| 9 | // the low end of our input first. |
| 10 | //Regardless of endianess, within any given byte the bits are always in most |
| 11 | // to least significant order. |
| 12 | //Also regardless of endianess, the buffer always aligns bits to the low end |
| 13 | // of the byte. |
| 14 | |
| 15 | /// Creates a bit writer which allows for writing bits to an underlying standard writer |
| 16 | pub fn BitWriter(comptime endian: std.builtin.Endian, comptime Writer: type) type { |
| 17 | return struct { |
| 18 | writer: Writer, |
| 19 | bits: u8 = 0, |
| 20 | count: u4 = 0, |
| 21 | |
| 22 | const low_bit_mask = [9]u8{ |
| 23 | 0b00000000, |
| 24 | 0b00000001, |
| 25 | 0b00000011, |
| 26 | 0b00000111, |
| 27 | 0b00001111, |
| 28 | 0b00011111, |
| 29 | 0b00111111, |
| 30 | 0b01111111, |
| 31 | 0b11111111, |
| 32 | }; |
| 33 | |
| 34 | /// Write the specified number of bits to the writer from the least significant bits of |
| 35 | /// the specified value. Bits will only be written to the writer when there |
| 36 | /// are enough to fill a byte. |
| 37 | pub fn writeBits(self: *@This(), value: anytype, num: u16) !void { |
| 38 | const T = @TypeOf(value); |
| 39 | const UT = std.meta.Int(.unsigned, @bitSizeOf(T)); |
| 40 | const U = if (@bitSizeOf(T) < 8) u8 else UT; //<u8 is a pain to work with |
| 41 | |
| 42 | var in: U = @as(UT, @bitCast(value)); |
| 43 | var in_count: u16 = num; |
| 44 | |
| 45 | if (self.count > 0) { |
| 46 | //if we can't fill the buffer, add what we have |
| 47 | const bits_free = 8 - self.count; |
| 48 | if (num < bits_free) { |
| 49 | self.addBits(@truncate(in), @intCast(num)); |
| 50 | return; |
| 51 | } |
| 52 | |
| 53 | //finish filling the buffer and flush it |
| 54 | if (num == bits_free) { |
| 55 | self.addBits(@truncate(in), @intCast(num)); |
| 56 | return self.flushBits(); |
| 57 | } |
| 58 | |
| 59 | switch (endian) { |
| 60 | .big => { |
| 61 | const bits = in >> @intCast(in_count - bits_free); |
| 62 | self.addBits(@truncate(bits), bits_free); |
| 63 | }, |
| 64 | .little => { |
| 65 | self.addBits(@truncate(in), bits_free); |
| 66 | in >>= @intCast(bits_free); |
| 67 | }, |
| 68 | } |
| 69 | in_count -= bits_free; |
| 70 | try self.flushBits(); |
| 71 | } |
| 72 | |
| 73 | //write full bytes while we can |
| 74 | const full_bytes_left = in_count / 8; |
| 75 | for (0..full_bytes_left) |_| { |
| 76 | switch (endian) { |
| 77 | .big => { |
| 78 | const bits = in >> @intCast(in_count - 8); |
| 79 | try self.writer.writeByte(@truncate(bits)); |
| 80 | }, |
| 81 | .little => { |
| 82 | try self.writer.writeByte(@truncate(in)); |
| 83 | if (U == u8) in = 0 else in >>= 8; |
| 84 | }, |
| 85 | } |
| 86 | in_count -= 8; |
| 87 | } |
| 88 | |
| 89 | //save the remaining bits in the buffer |
| 90 | self.addBits(@truncate(in), @intCast(in_count)); |
| 91 | } |
| 92 | |
| 93 | //convenience funciton for adding bits to the buffer |
| 94 | //in the appropriate position based on endianess |
| 95 | fn addBits(self: *@This(), bits: u8, num: u4) void { |
| 96 | if (num == 8) self.bits = bits else switch (endian) { |
| 97 | .big => { |
| 98 | self.bits <<= @intCast(num); |
| 99 | self.bits |= bits & low_bit_mask[num]; |
| 100 | }, |
| 101 | .little => { |
| 102 | const pos = bits << @intCast(self.count); |
| 103 | self.bits |= pos; |
| 104 | }, |
| 105 | } |
| 106 | self.count += num; |
| 107 | } |
| 108 | |
| 109 | /// Flush any remaining bits to the writer, filling |
| 110 | /// unused bits with 0s. |
| 111 | pub fn flushBits(self: *@This()) !void { |
| 112 | if (self.count == 0) return; |
| 113 | if (endian == .big) self.bits <<= @intCast(8 - self.count); |
| 114 | try self.writer.writeByte(self.bits); |
| 115 | self.bits = 0; |
| 116 | self.count = 0; |
| 117 | } |
| 118 | }; |
| 119 | } |
| 120 | |
| 121 | pub fn bitWriter(comptime endian: std.builtin.Endian, writer: anytype) BitWriter(endian, @TypeOf(writer)) { |
| 122 | return .{ .writer = writer }; |
| 123 | } |
| 124 | |
| 125 | /////////////////////////////// |
| 126 | |
| 127 | test "api coverage" { |
| 128 | var mem_be = [_]u8{0} ** 2; |
| 129 | var mem_le = [_]u8{0} ** 2; |
| 130 | |
| 131 | var mem_out_be = std.io.fixedBufferStream(&mem_be); |
| 132 | var bit_stream_be = bitWriter(.big, mem_out_be.writer()); |
| 133 | |
| 134 | const testing = std.testing; |
| 135 | |
| 136 | try bit_stream_be.writeBits(@as(u2, 1), 1); |
| 137 | try bit_stream_be.writeBits(@as(u5, 2), 2); |
| 138 | try bit_stream_be.writeBits(@as(u128, 3), 3); |
| 139 | try bit_stream_be.writeBits(@as(u8, 4), 4); |
| 140 | try bit_stream_be.writeBits(@as(u9, 5), 5); |
| 141 | try bit_stream_be.writeBits(@as(u1, 1), 1); |
| 142 | |
| 143 | try testing.expect(mem_be[0] == 0b11001101 and mem_be[1] == 0b00001011); |
| 144 | |
| 145 | mem_out_be.pos = 0; |
| 146 | |
| 147 | try bit_stream_be.writeBits(@as(u15, 0b110011010000101), 15); |
| 148 | try bit_stream_be.flushBits(); |
| 149 | try testing.expect(mem_be[0] == 0b11001101 and mem_be[1] == 0b00001010); |
| 150 | |
| 151 | mem_out_be.pos = 0; |
| 152 | try bit_stream_be.writeBits(@as(u32, 0b110011010000101), 16); |
| 153 | try testing.expect(mem_be[0] == 0b01100110 and mem_be[1] == 0b10000101); |
| 154 | |
| 155 | try bit_stream_be.writeBits(@as(u0, 0), 0); |
| 156 | |
| 157 | var mem_out_le = std.io.fixedBufferStream(&mem_le); |
| 158 | var bit_stream_le = bitWriter(.little, mem_out_le.writer()); |
| 159 | |
| 160 | try bit_stream_le.writeBits(@as(u2, 1), 1); |
| 161 | try bit_stream_le.writeBits(@as(u5, 2), 2); |
| 162 | try bit_stream_le.writeBits(@as(u128, 3), 3); |
| 163 | try bit_stream_le.writeBits(@as(u8, 4), 4); |
| 164 | try bit_stream_le.writeBits(@as(u9, 5), 5); |
| 165 | try bit_stream_le.writeBits(@as(u1, 1), 1); |
| 166 | |
| 167 | try testing.expect(mem_le[0] == 0b00011101 and mem_le[1] == 0b10010101); |
| 168 | |
| 169 | mem_out_le.pos = 0; |
| 170 | try bit_stream_le.writeBits(@as(u15, 0b110011010000101), 15); |
| 171 | try bit_stream_le.flushBits(); |
| 172 | try testing.expect(mem_le[0] == 0b10000101 and mem_le[1] == 0b01100110); |
| 173 | |
| 174 | mem_out_le.pos = 0; |
| 175 | try bit_stream_le.writeBits(@as(u32, 0b1100110100001011), 16); |
| 176 | try testing.expect(mem_le[0] == 0b00001011 and mem_le[1] == 0b11001101); |
| 177 | |
| 178 | try bit_stream_le.writeBits(@as(u0, 0), 0); |
| 179 | } |
Andrew Kelley <andrew@ziglang.org> 2025-05-06 20:25:55-07:00