authorgravatar for tgschultz@gmail.comtgschultz <tgschultz@gmail.com> 2024-10-13 20:44:42-05:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2024-10-13 18:44:42-07:00
logba569bb8e9c3f57e5159073c5b92d5bcdcfaf076
tree37416ae936343c6ed08658d63fa5d34132a51929
parente2e79960d2067f25fa66945e73c6cd81cad39804
signaturebadge-check Signed by PGP key B5690EEEBB952194

Rewrite bit_reader and bit_writer to take advantage of current zig semantics and enhance readability (#21689)

Co-authored-by: Tanner Schultz <tgschultz@tgschultz-dl.tail7ba92.ts.net>

6 files changed, 244 insertions(+), 256 deletions(-)

lib/std/compress/zstandard/decode/block.zig+1-1
...@@ -405,7 +405,7 @@ pub const DecodeState = struct {...@@ -405,7 +405,7 @@ pub const DecodeState = struct {
405 };405 };
406 fn readLiteralsBits(406 fn readLiteralsBits(
407 self: *DecodeState,407 self: *DecodeState,
408 bit_count_to_read: usize,408 bit_count_to_read: u16,
409 ) LiteralBitsError!u16 {409 ) LiteralBitsError!u16 {
410 return self.literal_stream_reader.readBitsNoEof(u16, bit_count_to_read) catch bits: {410 return self.literal_stream_reader.readBitsNoEof(u16, bit_count_to_read) catch bits: {
411 if (self.literal_streams == .four and self.literal_stream_index < 3) {411 if (self.literal_streams == .four and self.literal_stream_index < 3) {
lib/std/compress/zstandard/decode/huffman.zig+2-2
...@@ -63,7 +63,7 @@ fn decodeFseHuffmanTreeSlice(src: []const u8, compressed_size: usize, weights: *...@@ -63,7 +63,7 @@ fn decodeFseHuffmanTreeSlice(src: []const u8, compressed_size: usize, weights: *
6363
64fn assignWeights(64fn assignWeights(
65 huff_bits: *readers.ReverseBitReader,65 huff_bits: *readers.ReverseBitReader,
66 accuracy_log: usize,66 accuracy_log: u16,
67 entries: *[1 << 6]Table.Fse,67 entries: *[1 << 6]Table.Fse,
68 weights: *[256]u4,68 weights: *[256]u4,
69) !usize {69) !usize {
...@@ -73,7 +73,7 @@ fn assignWeights(...@@ -73,7 +73,7 @@ fn assignWeights(
7373
74 while (i < 254) {74 while (i < 254) {
75 const even_data = entries[even_state];75 const even_data = entries[even_state];
76 var read_bits: usize = 0;76 var read_bits: u16 = 0;
77 const even_bits = huff_bits.readBits(u32, even_data.bits, &read_bits) catch unreachable;77 const even_bits = huff_bits.readBits(u32, even_data.bits, &read_bits) catch unreachable;
78 weights[i] = std.math.cast(u4, even_data.symbol) orelse return error.MalformedHuffmanTree;78 weights[i] = std.math.cast(u4, even_data.symbol) orelse return error.MalformedHuffmanTree;
79 i += 1;79 i += 1;
lib/std/compress/zstandard/readers.zig+5-5
...@@ -42,11 +42,11 @@ pub const ReverseBitReader = struct {...@@ -42,11 +42,11 @@ pub const ReverseBitReader = struct {
42 if (i == 8) return error.BitStreamHasNoStartBit;42 if (i == 8) return error.BitStreamHasNoStartBit;
43 }43 }
4444
45 pub fn readBitsNoEof(self: *@This(), comptime U: type, num_bits: usize) error{EndOfStream}!U {45 pub fn readBitsNoEof(self: *@This(), comptime U: type, num_bits: u16) error{EndOfStream}!U {
46 return self.bit_reader.readBitsNoEof(U, num_bits);46 return self.bit_reader.readBitsNoEof(U, num_bits);
47 }47 }
4848
49 pub fn readBits(self: *@This(), comptime U: type, num_bits: usize, out_bits: *usize) error{}!U {49 pub fn readBits(self: *@This(), comptime U: type, num_bits: u16, out_bits: *u16) error{}!U {
50 return try self.bit_reader.readBits(U, num_bits, out_bits);50 return try self.bit_reader.readBits(U, num_bits, out_bits);
51 }51 }
5252
...@@ -55,7 +55,7 @@ pub const ReverseBitReader = struct {...@@ -55,7 +55,7 @@ pub const ReverseBitReader = struct {
55 }55 }
5656
57 pub fn isEmpty(self: ReverseBitReader) bool {57 pub fn isEmpty(self: ReverseBitReader) bool {
58 return self.byte_reader.remaining_bytes == 0 and self.bit_reader.bit_count == 0;58 return self.byte_reader.remaining_bytes == 0 and self.bit_reader.count == 0;
59 }59 }
60};60};
6161
...@@ -63,11 +63,11 @@ pub fn BitReader(comptime Reader: type) type {...@@ -63,11 +63,11 @@ pub fn BitReader(comptime Reader: type) type {
63 return struct {63 return struct {
64 underlying: std.io.BitReader(.little, Reader),64 underlying: std.io.BitReader(.little, Reader),
6565
66 pub fn readBitsNoEof(self: *@This(), comptime U: type, num_bits: usize) !U {66 pub fn readBitsNoEof(self: *@This(), comptime U: type, num_bits: u16) !U {
67 return self.underlying.readBitsNoEof(U, num_bits);67 return self.underlying.readBitsNoEof(U, num_bits);
68 }68 }
6969
70 pub fn readBits(self: *@This(), comptime U: type, num_bits: usize, out_bits: *usize) !U {70 pub fn readBits(self: *@This(), comptime U: type, num_bits: u16, out_bits: *u16) !U {
71 return self.underlying.readBits(U, num_bits, out_bits);71 return self.underlying.readBits(U, num_bits, out_bits);
72 }72 }
7373
lib/std/io/bit_reader.zig+136-133
...@@ -1,176 +1,179 @@...@@ -1,176 +1,179 @@
1const std = @import("../std.zig");1const std = @import("../std.zig");
2const io = std.io;2
3const assert = std.debug.assert;3//General note on endianess:
4const testing = std.testing;4//Big endian is packed starting in the most significant part of the byte and subsequent
5const meta = std.meta;5// bytes contain less significant bits. Thus we always take bits from the high
6const math = std.math;6// end and place them below existing bits in our output.
77//Little endian is packed starting in the least significant part of the byte and
8/// Creates a stream which allows for reading bit fields from another stream8// subsequent bytes contain more significant bits. Thus we always take bits from
9pub fn BitReader(comptime endian: std.builtin.Endian, comptime ReaderType: type) type {9// the low end and place them above existing bits in our output.
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 reader which allows for reading bits from an underlying standard reader
16pub fn BitReader(comptime endian: std.builtin.Endian, comptime Reader: type) type {
10 return struct {17 return struct {
11 forward_reader: ReaderType,18 reader: Reader,
12 bit_buffer: u7,19 bits: u8 = 0,
13 bit_count: u3,20 count: u4 = 0,
1421
15 pub const Error = ReaderType.Error;22 const low_bit_mask = [9]u8{
16 pub const Reader = io.Reader(*Self, Error, read);23 0b00000000,
1724 0b00000001,
18 const Self = @This();25 0b00000011,
19 const u8_bit_count = @bitSizeOf(u8);26 0b00000111,
20 const u7_bit_count = @bitSizeOf(u7);27 0b00001111,
21 const u4_bit_count = @bitSizeOf(u4);28 0b00011111,
2229 0b00111111,
23 pub fn init(forward_reader: ReaderType) Self {30 0b01111111,
24 return Self{31 0b11111111,
25 .forward_reader = forward_reader,32 };
26 .bit_buffer = 0,33
27 .bit_count = 0,34 fn Bits(comptime T: type) type {
35 return struct {
36 T,
37 u16,
38 };
39 }
40
41 fn initBits(comptime T: type, out: anytype, num: u16) Bits(T) {
42 const UT = std.meta.Int(.unsigned, @bitSizeOf(T));
43 return .{
44 @bitCast(@as(UT, @intCast(out))),
45 num,
28 };46 };
29 }47 }
3048
31 /// Reads `bits` bits from the stream and returns a specified unsigned int type49 /// Reads `bits` bits from the reader and returns a specified type
32 /// containing them in the least significant end, returning an error if the50 /// containing them in the least significant end, returning an error if the
33 /// specified number of bits could not be read.51 /// specified number of bits could not be read.
34 pub fn readBitsNoEof(self: *Self, comptime U: type, bits: usize) !U {52 pub fn readBitsNoEof(self: *@This(), comptime T: type, num: u16) !T {
35 var n: usize = undefined;53 const b, const c = try self.readBitsTuple(T, num);
36 const result = try self.readBits(U, bits, &n);54 if (c < num) return error.EndOfStream;
37 if (n < bits) return error.EndOfStream;55 return b;
38 return result;
39 }56 }
4057
41 /// Reads `bits` bits from the stream and returns a specified unsigned int type58 /// Reads `bits` bits from the reader and returns a specified type
42 /// containing them in the least significant end. The number of bits successfully59 /// containing them in the least significant end. The number of bits successfully
43 /// read is placed in `out_bits`, as reaching the end of the stream is not an error.60 /// read is placed in `out_bits`, as reaching the end of the stream is not an error.
44 pub fn readBits(self: *Self, comptime U: type, bits: usize, out_bits: *usize) Error!U {61 pub fn readBits(self: *@This(), comptime T: type, num: u16, out_bits: *u16) !T {
45 //by extending the buffer to a minimum of u8 we can cover a number of edge cases62 const b, const c = try self.readBitsTuple(T, num);
46 // related to shifting and casting.63 out_bits.* = c;
47 const u_bit_count = @bitSizeOf(U);64 return b;
48 const buf_bit_count = bc: {65 }
49 assert(u_bit_count >= bits);
50 break :bc if (u_bit_count <= u8_bit_count) u8_bit_count else u_bit_count;
51 };
52 const Buf = std.meta.Int(.unsigned, buf_bit_count);
53 const BufShift = math.Log2Int(Buf);
5466
55 out_bits.* = @as(usize, 0);67 /// Reads `bits` bits from the reader and returns a tuple of the specified type
56 if (U == u0 or bits == 0) return 0;68 /// containing them in the least significant end, and the number of bits successfully
57 var out_buffer = @as(Buf, 0);69 /// read. Reaching the end of the stream is not an error.
70 pub fn readBitsTuple(self: *@This(), comptime T: type, num: u16) !Bits(T) {
71 const UT = std.meta.Int(.unsigned, @bitSizeOf(T));
72 const U = if (@bitSizeOf(T) < 8) u8 else UT; //it is a pain to work with <u8
5873
59 if (self.bit_count > 0) {74 //dump any bits in our buffer first
60 const n = if (self.bit_count >= bits) @as(u3, @intCast(bits)) else self.bit_count;75 if (num <= self.count) return initBits(T, self.removeBits(@intCast(num)), num);
61 const shift = u7_bit_count - n;
62 switch (endian) {
63 .big => {
64 out_buffer = @as(Buf, self.bit_buffer >> shift);
65 if (n >= u7_bit_count)
66 self.bit_buffer = 0
67 else
68 self.bit_buffer <<= n;
69 },
70 .little => {
71 const value = (self.bit_buffer << shift) >> shift;
72 out_buffer = @as(Buf, value);
73 if (n >= u7_bit_count)
74 self.bit_buffer = 0
75 else
76 self.bit_buffer >>= n;
77 },
78 }
79 self.bit_count -= n;
80 out_bits.* = n;
81 }
82 //at this point we know bit_buffer is empty
8376
84 //copy bytes until we have enough bits, then leave the rest in bit_buffer77 var out_count: u16 = self.count;
85 while (out_bits.* < bits) {78 var out: U = self.removeBits(self.count);
86 const n = bits - out_bits.*;79
87 const next_byte = self.forward_reader.readByte() catch |err| switch (err) {80 //grab all the full bytes we need and put their
88 error.EndOfStream => return @as(U, @intCast(out_buffer)),81 //bits where they belong
82 const full_bytes_left = (num - out_count) / 8;
83
84 for (0..full_bytes_left) |_| {
85 const byte = self.reader.readByte() catch |err| switch (err) {
86 error.EndOfStream => return initBits(T, out, out_count),
89 else => |e| return e,87 else => |e| return e,
90 };88 };
9189
92 switch (endian) {90 switch (endian) {
93 .big => {91 .big => {
94 if (n >= u8_bit_count) {92 if (U == u8) out = 0 else out <<= 8; //shifting u8 by 8 is illegal in Zig
95 out_buffer <<= @as(u3, @intCast(u8_bit_count - 1));93 out |= byte;
96 out_buffer <<= 1;
97 out_buffer |= @as(Buf, next_byte);
98 out_bits.* += u8_bit_count;
99 continue;
100 }
101
102 const shift = @as(u3, @intCast(u8_bit_count - n));
103 out_buffer <<= @as(BufShift, @intCast(n));
104 out_buffer |= @as(Buf, next_byte >> shift);
105 out_bits.* += n;
106 self.bit_buffer = @as(u7, @truncate(next_byte << @as(u3, @intCast(n - 1))));
107 self.bit_count = shift;
108 },94 },
109 .little => {95 .little => {
110 if (n >= u8_bit_count) {96 const pos = @as(U, byte) << @intCast(out_count);
111 out_buffer |= @as(Buf, next_byte) << @as(BufShift, @intCast(out_bits.*));97 out |= pos;
112 out_bits.* += u8_bit_count;
113 continue;
114 }
115
116 const shift = @as(u3, @intCast(u8_bit_count - n));
117 const value = (next_byte << shift) >> shift;
118 out_buffer |= @as(Buf, value) << @as(BufShift, @intCast(out_bits.*));
119 out_bits.* += n;
120 self.bit_buffer = @as(u7, @truncate(next_byte >> @as(u3, @intCast(n))));
121 self.bit_count = shift;
122 },98 },
123 }99 }
100 out_count += 8;
124 }101 }
125102
126 return @as(U, @intCast(out_buffer));103 const bits_left = num - out_count;
127 }104 const keep = 8 - bits_left;
105
106 if (bits_left == 0) return initBits(T, out, out_count);
128107
129 pub fn alignToByte(self: *Self) void {108 const final_byte = self.reader.readByte() catch |err| switch (err) {
130 self.bit_buffer = 0;109 error.EndOfStream => return initBits(T, out, out_count),
131 self.bit_count = 0;110 else => |e| return e,
111 };
112
113 switch (endian) {
114 .big => {
115 out <<= @intCast(bits_left);
116 out |= final_byte >> @intCast(keep);
117 self.bits = final_byte & low_bit_mask[keep];
118 },
119 .little => {
120 const pos = @as(U, final_byte & low_bit_mask[bits_left]) << @intCast(out_count);
121 out |= pos;
122 self.bits = final_byte >> @intCast(bits_left);
123 },
124 }
125
126 self.count = @intCast(keep);
127 return initBits(T, out, num);
132 }128 }
133129
134 pub fn read(self: *Self, buffer: []u8) Error!usize {130 //convenience function for removing bits from
135 var out_bits: usize = undefined;131 //the appropriate part of the buffer based on
136 var out_bits_total = @as(usize, 0);132 //endianess.
137 //@NOTE: I'm not sure this is a good idea, maybe alignToByte should be forced133 fn removeBits(self: *@This(), num: u4) u8 {
138 if (self.bit_count > 0) {134 if (num == 8) {
139 for (buffer) |*b| {135 self.count = 0;
140 b.* = try self.readBits(u8, u8_bit_count, &out_bits);136 return self.bits;
141 out_bits_total += out_bits;137 }
142 }138
143 const incomplete_byte = @intFromBool(out_bits_total % u8_bit_count > 0);139 const keep = self.count - num;
144 return (out_bits_total / u8_bit_count) + incomplete_byte;140 const bits = switch (endian) {
141 .big => self.bits >> @intCast(keep),
142 .little => self.bits & low_bit_mask[num],
143 };
144 switch (endian) {
145 .big => self.bits &= low_bit_mask[keep],
146 .little => self.bits >>= @intCast(num),
145 }147 }
146148
147 return self.forward_reader.read(buffer);149 self.count = keep;
150 return bits;
148 }151 }
149152
150 pub fn reader(self: *Self) Reader {153 pub fn alignToByte(self: *@This()) void {
151 return .{ .context = self };154 self.bits = 0;
155 self.count = 0;
152 }156 }
153 };157 };
154}158}
155159
156pub fn bitReader(160pub fn bitReader(comptime endian: std.builtin.Endian, reader: anytype) BitReader(endian, @TypeOf(reader)) {
157 comptime endian: std.builtin.Endian,161 return .{ .reader = reader };
158 underlying_stream: anytype,
159) BitReader(endian, @TypeOf(underlying_stream)) {
160 return BitReader(endian, @TypeOf(underlying_stream)).init(underlying_stream);
161}162}
162163
164///////////////////////////////
165
163test "api coverage" {166test "api coverage" {
164 const mem_be = [_]u8{ 0b11001101, 0b00001011 };167 const mem_be = [_]u8{ 0b11001101, 0b00001011 };
165 const mem_le = [_]u8{ 0b00011101, 0b10010101 };168 const mem_le = [_]u8{ 0b00011101, 0b10010101 };
166169
167 var mem_in_be = io.fixedBufferStream(&mem_be);170 var mem_in_be = std.io.fixedBufferStream(&mem_be);
168 var bit_stream_be = bitReader(.big, mem_in_be.reader());171 var bit_stream_be = bitReader(.big, mem_in_be.reader());
169172
170 var out_bits: usize = undefined;173 var out_bits: u16 = undefined;
171174
172 const expect = testing.expect;175 const expect = std.testing.expect;
173 const expectError = testing.expectError;176 const expectError = std.testing.expectError;
174177
175 try expect(1 == try bit_stream_be.readBits(u2, 1, &out_bits));178 try expect(1 == try bit_stream_be.readBits(u2, 1, &out_bits));
176 try expect(out_bits == 1);179 try expect(out_bits == 1);
...@@ -186,12 +189,12 @@ test "api coverage" {...@@ -186,12 +189,12 @@ test "api coverage" {
186 try expect(out_bits == 1);189 try expect(out_bits == 1);
187190
188 mem_in_be.pos = 0;191 mem_in_be.pos = 0;
189 bit_stream_be.bit_count = 0;192 bit_stream_be.count = 0;
190 try expect(0b110011010000101 == try bit_stream_be.readBits(u15, 15, &out_bits));193 try expect(0b110011010000101 == try bit_stream_be.readBits(u15, 15, &out_bits));
191 try expect(out_bits == 15);194 try expect(out_bits == 15);
192195
193 mem_in_be.pos = 0;196 mem_in_be.pos = 0;
194 bit_stream_be.bit_count = 0;197 bit_stream_be.count = 0;
195 try expect(0b1100110100001011 == try bit_stream_be.readBits(u16, 16, &out_bits));198 try expect(0b1100110100001011 == try bit_stream_be.readBits(u16, 16, &out_bits));
196 try expect(out_bits == 16);199 try expect(out_bits == 16);
197200
...@@ -201,7 +204,7 @@ test "api coverage" {...@@ -201,7 +204,7 @@ test "api coverage" {
201 try expect(out_bits == 0);204 try expect(out_bits == 0);
202 try expectError(error.EndOfStream, bit_stream_be.readBitsNoEof(u1, 1));205 try expectError(error.EndOfStream, bit_stream_be.readBitsNoEof(u1, 1));
203206
204 var mem_in_le = io.fixedBufferStream(&mem_le);207 var mem_in_le = std.io.fixedBufferStream(&mem_le);
205 var bit_stream_le = bitReader(.little, mem_in_le.reader());208 var bit_stream_le = bitReader(.little, mem_in_le.reader());
206209
207 try expect(1 == try bit_stream_le.readBits(u2, 1, &out_bits));210 try expect(1 == try bit_stream_le.readBits(u2, 1, &out_bits));
...@@ -218,12 +221,12 @@ test "api coverage" {...@@ -218,12 +221,12 @@ test "api coverage" {
218 try expect(out_bits == 1);221 try expect(out_bits == 1);
219222
220 mem_in_le.pos = 0;223 mem_in_le.pos = 0;
221 bit_stream_le.bit_count = 0;224 bit_stream_le.count = 0;
222 try expect(0b001010100011101 == try bit_stream_le.readBits(u15, 15, &out_bits));225 try expect(0b001010100011101 == try bit_stream_le.readBits(u15, 15, &out_bits));
223 try expect(out_bits == 15);226 try expect(out_bits == 15);
224227
225 mem_in_le.pos = 0;228 mem_in_le.pos = 0;
226 bit_stream_le.bit_count = 0;229 bit_stream_le.count = 0;
227 try expect(0b1001010100011101 == try bit_stream_le.readBits(u16, 16, &out_bits));230 try expect(0b1001010100011101 == try bit_stream_le.readBits(u16, 16, &out_bits));
228 try expect(out_bits == 16);231 try expect(out_bits == 16);
229232
lib/std/io/bit_writer.zig+99-114
...@@ -1,153 +1,138 @@...@@ -1,153 +1,138 @@
1const std = @import("../std.zig");1const std = @import("../std.zig");
2const io = std.io;
3const testing = std.testing;
4const assert = std.debug.assert;
5const math = std.math;
62
7/// Creates a stream which allows for writing bit fields to another stream3//General note on endianess:
8pub fn BitWriter(comptime endian: std.builtin.Endian, comptime WriterType: type) type {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
16pub fn BitWriter(comptime endian: std.builtin.Endian, comptime Writer: type) type {
9 return struct {17 return struct {
10 forward_writer: WriterType,18 writer: Writer,
11 bit_buffer: u8,19 bits: u8 = 0,
12 bit_count: u4,20 count: u4 = 0,
1321
14 pub const Error = WriterType.Error;22 const low_bit_mask = [9]u8{
15 pub const Writer = io.Writer(*Self, Error, write);23 0b00000000,
1624 0b00000001,
17 const Self = @This();25 0b00000011,
18 const u8_bit_count = @bitSizeOf(u8);26 0b00000111,
19 const u4_bit_count = @bitSizeOf(u4);27 0b00001111,
2028 0b00011111,
21 pub fn init(forward_writer: WriterType) Self {29 0b00111111,
22 return Self{30 0b01111111,
23 .forward_writer = forward_writer,31 0b11111111,
24 .bit_buffer = 0,32 };
25 .bit_count = 0,33
26 };34 /// Write the specified number of bits to the writer from the least significant bits of
27 }35 /// the specified value. Bits will only be written to the writer when there
28
29 /// Write the specified number of bits to the stream from the least significant bits of
30 /// the specified unsigned int value. Bits will only be written to the stream when there
31 /// are enough to fill a byte.36 /// are enough to fill a byte.
32 pub fn writeBits(self: *Self, value: anytype, bits: usize) Error!void {37 pub fn writeBits(self: *@This(), value: anytype, num: u16) !void {
33 if (bits == 0) return;38 const T = @TypeOf(value);
3439 const UT = std.meta.Int(.unsigned, @bitSizeOf(T));
35 const U = @TypeOf(value);40 const U = if (@bitSizeOf(T) < 8) u8 else UT; //<u8 is a pain to work with
36 comptime assert(@typeInfo(U).int.signedness == .unsigned);41
3742 var in: U = @as(UT, @bitCast(value));
38 //by extending the buffer to a minimum of u8 we can cover a number of edge cases43 var in_count: u16 = num;
39 // related to shifting and casting.44
40 const u_bit_count = @bitSizeOf(U);45 if (self.count > 0) {
41 const buf_bit_count = bc: {46 //if we can't fill the buffer, add what we have
42 assert(u_bit_count >= bits);47 const bits_free = 8 - self.count;
43 break :bc if (u_bit_count <= u8_bit_count) u8_bit_count else u_bit_count;48 if (num < bits_free) {
44 };49 self.addBits(@truncate(in), @intCast(num));
45 const Buf = std.meta.Int(.unsigned, buf_bit_count);50 return;
46 const BufShift = math.Log2Int(Buf);51 }
4752
48 const buf_value = @as(Buf, @intCast(value));53 //finish filling the buffer and flush it
4954 if (num == bits_free) {
50 const high_byte_shift = @as(BufShift, @intCast(buf_bit_count - u8_bit_count));55 self.addBits(@truncate(in), @intCast(num));
51 var in_buffer = switch (endian) {56 return self.flushBits();
52 .big => buf_value << @as(BufShift, @intCast(buf_bit_count - bits)),57 }
53 .little => buf_value,58
54 };
55 var in_bits = bits;
56
57 if (self.bit_count > 0) {
58 const bits_remaining = u8_bit_count - self.bit_count;
59 const n = @as(u3, @intCast(if (bits_remaining > bits) bits else bits_remaining));
60 switch (endian) {59 switch (endian) {
61 .big => {60 .big => {
62 const shift = @as(BufShift, @intCast(high_byte_shift + self.bit_count));61 const bits = in >> @intCast(in_count - bits_free);
63 const v = @as(u8, @intCast(in_buffer >> shift));62 self.addBits(@truncate(bits), bits_free);
64 self.bit_buffer |= v;
65 in_buffer <<= n;
66 },63 },
67 .little => {64 .little => {
68 const v = @as(u8, @truncate(in_buffer)) << @as(u3, @intCast(self.bit_count));65 self.addBits(@truncate(in), bits_free);
69 self.bit_buffer |= v;66 in >>= @intCast(bits_free);
70 in_buffer >>= n;
71 },67 },
72 }68 }
73 self.bit_count += n;69 in_count -= bits_free;
74 in_bits -= n;70 try self.flushBits();
75
76 //if we didn't fill the buffer, it's because bits < bits_remaining;
77 if (self.bit_count != u8_bit_count) return;
78 try self.forward_writer.writeByte(self.bit_buffer);
79 self.bit_buffer = 0;
80 self.bit_count = 0;
81 }71 }
82 //at this point we know bit_buffer is empty
8372
84 //copy bytes until we can't fill one anymore, then leave the rest in bit_buffer73 //write full bytes while we can
85 while (in_bits >= u8_bit_count) {74 const full_bytes_left = in_count / 8;
75 for (0..full_bytes_left) |_| {
86 switch (endian) {76 switch (endian) {
87 .big => {77 .big => {
88 const v = @as(u8, @intCast(in_buffer >> high_byte_shift));78 const bits = in >> @intCast(in_count - 8);
89 try self.forward_writer.writeByte(v);79 try self.writer.writeByte(@truncate(bits));
90 in_buffer <<= @as(u3, @intCast(u8_bit_count - 1));
91 in_buffer <<= 1;
92 },80 },
93 .little => {81 .little => {
94 const v = @as(u8, @truncate(in_buffer));82 try self.writer.writeByte(@truncate(in));
95 try self.forward_writer.writeByte(v);83 if (U == u8) in = 0 else in >>= 8;
96 in_buffer >>= @as(u3, @intCast(u8_bit_count - 1));
97 in_buffer >>= 1;
98 },84 },
99 }85 }
100 in_bits -= u8_bit_count;86 in_count -= 8;
101 }87 }
10288
103 if (in_bits > 0) {89 //save the remaining bits in the buffer
104 self.bit_count = @as(u4, @intCast(in_bits));90 self.addBits(@truncate(in), @intCast(in_count));
105 self.bit_buffer = switch (endian) {
106 .big => @as(u8, @truncate(in_buffer >> high_byte_shift)),
107 .little => @as(u8, @truncate(in_buffer)),
108 };
109 }
110 }
111
112 /// Flush any remaining bits to the stream.
113 pub fn flushBits(self: *Self) Error!void {
114 if (self.bit_count == 0) return;
115 try self.forward_writer.writeByte(self.bit_buffer);
116 self.bit_buffer = 0;
117 self.bit_count = 0;
118 }91 }
11992
120 pub fn write(self: *Self, buffer: []const u8) Error!usize {93 //convenience funciton for adding bits to the buffer
121 // TODO: I'm not sure this is a good idea, maybe flushBits should be forced94 //in the appropriate position based on endianess
122 if (self.bit_count > 0) {95 fn addBits(self: *@This(), bits: u8, num: u4) void {
123 for (buffer) |b|96 if (num == 8) self.bits = bits else switch (endian) {
124 try self.writeBits(b, u8_bit_count);97 .big => {
125 return buffer.len;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 },
126 }105 }
127106 self.count += num;
128 return self.forward_writer.write(buffer);
129 }107 }
130108
131 pub fn writer(self: *Self) Writer {109 /// Flush any remaining bits to the writer, filling
132 return .{ .context = self };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;
133 }117 }
134 };118 };
135}119}
136120
137pub fn bitWriter(121pub fn bitWriter(comptime endian: std.builtin.Endian, writer: anytype) BitWriter(endian, @TypeOf(writer)) {
138 comptime endian: std.builtin.Endian,122 return .{ .writer = writer };
139 underlying_stream: anytype,
140) BitWriter(endian, @TypeOf(underlying_stream)) {
141 return BitWriter(endian, @TypeOf(underlying_stream)).init(underlying_stream);
142}123}
143124
125///////////////////////////////
126
144test "api coverage" {127test "api coverage" {
145 var mem_be = [_]u8{0} ** 2;128 var mem_be = [_]u8{0} ** 2;
146 var mem_le = [_]u8{0} ** 2;129 var mem_le = [_]u8{0} ** 2;
147130
148 var mem_out_be = io.fixedBufferStream(&mem_be);131 var mem_out_be = std.io.fixedBufferStream(&mem_be);
149 var bit_stream_be = bitWriter(.big, mem_out_be.writer());132 var bit_stream_be = bitWriter(.big, mem_out_be.writer());
150133
134 const testing = std.testing;
135
151 try bit_stream_be.writeBits(@as(u2, 1), 1);136 try bit_stream_be.writeBits(@as(u2, 1), 1);
152 try bit_stream_be.writeBits(@as(u5, 2), 2);137 try bit_stream_be.writeBits(@as(u5, 2), 2);
153 try bit_stream_be.writeBits(@as(u128, 3), 3);138 try bit_stream_be.writeBits(@as(u128, 3), 3);
...@@ -169,7 +154,7 @@ test "api coverage" {...@@ -169,7 +154,7 @@ test "api coverage" {
169154
170 try bit_stream_be.writeBits(@as(u0, 0), 0);155 try bit_stream_be.writeBits(@as(u0, 0), 0);
171156
172 var mem_out_le = io.fixedBufferStream(&mem_le);157 var mem_out_le = std.io.fixedBufferStream(&mem_le);
173 var bit_stream_le = bitWriter(.little, mem_out_le.writer());158 var bit_stream_le = bitWriter(.little, mem_out_le.writer());
174159
175 try bit_stream_le.writeBits(@as(u2, 1), 1);160 try bit_stream_le.writeBits(@as(u2, 1), 1);
lib/std/io/test.zig+1-1
...@@ -82,7 +82,7 @@ test "BitStreams with File Stream" {...@@ -82,7 +82,7 @@ test "BitStreams with File Stream" {
8282
83 var bit_stream = io.bitReader(native_endian, file.reader());83 var bit_stream = io.bitReader(native_endian, file.reader());
8484
85 var out_bits: usize = undefined;85 var out_bits: u16 = undefined;
8686
87 try expect(1 == try bit_stream.readBits(u2, 1, &out_bits));87 try expect(1 == try bit_stream.readBits(u2, 1, &out_bits));
88 try expect(out_bits == 1);88 try expect(out_bits == 1);