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 {
405405 };
406406 fn readLiteralsBits(
407407 self: *DecodeState,
408 bit_count_to_read: usize,
408 bit_count_to_read: u16,
409409 ) LiteralBitsError!u16 {
410410 return self.literal_stream_reader.readBitsNoEof(u16, bit_count_to_read) catch bits: {
411411 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: *
6363
6464fn assignWeights(
6565 huff_bits: *readers.ReverseBitReader,
66 accuracy_log: usize,
66 accuracy_log: u16,
6767 entries: *[1 << 6]Table.Fse,
6868 weights: *[256]u4,
6969) !usize {
......@@ -73,7 +73,7 @@ fn assignWeights(
7373
7474 while (i < 254) {
7575 const even_data = entries[even_state];
76 var read_bits: usize = 0;
76 var read_bits: u16 = 0;
7777 const even_bits = huff_bits.readBits(u32, even_data.bits, &read_bits) catch unreachable;
7878 weights[i] = std.math.cast(u4, even_data.symbol) orelse return error.MalformedHuffmanTree;
7979 i += 1;
lib/std/compress/zstandard/readers.zig+5-5
......@@ -42,11 +42,11 @@ pub const ReverseBitReader = struct {
4242 if (i == 8) return error.BitStreamHasNoStartBit;
4343 }
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 {
4646 return self.bit_reader.readBitsNoEof(U, num_bits);
4747 }
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 {
5050 return try self.bit_reader.readBits(U, num_bits, out_bits);
5151 }
5252
......@@ -55,7 +55,7 @@ pub const ReverseBitReader = struct {
5555 }
5656
5757 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;
5959 }
6060};
6161
......@@ -63,11 +63,11 @@ pub fn BitReader(comptime Reader: type) type {
6363 return struct {
6464 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 {
6767 return self.underlying.readBitsNoEof(U, num_bits);
6868 }
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 {
7171 return self.underlying.readBits(U, num_bits, out_bits);
7272 }
7373
lib/std/io/bit_reader.zig+136-133
......@@ -1,176 +1,179 @@
11const std = @import("../std.zig");
2const io = std.io;
3const assert = std.debug.assert;
4const testing = std.testing;
5const meta = std.meta;
6const math = std.math;
7
8/// Creates a stream which allows for reading bit fields from another stream
9pub fn BitReader(comptime endian: std.builtin.Endian, comptime ReaderType: type) type {
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 always take bits from the high
6// end and place them below existing bits in our output.
7//Little endian is packed starting in the least significant part of the byte and
8// subsequent bytes contain more significant bits. Thus we always take bits from
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 {
1017 return struct {
11 forward_reader: ReaderType,
12 bit_buffer: u7,
13 bit_count: u3,
14
15 pub const Error = ReaderType.Error;
16 pub const Reader = io.Reader(*Self, Error, read);
17
18 const Self = @This();
19 const u8_bit_count = @bitSizeOf(u8);
20 const u7_bit_count = @bitSizeOf(u7);
21 const u4_bit_count = @bitSizeOf(u4);
22
23 pub fn init(forward_reader: ReaderType) Self {
24 return Self{
25 .forward_reader = forward_reader,
26 .bit_buffer = 0,
27 .bit_count = 0,
18 reader: Reader,
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 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,
2846 };
2947 }
3048
31 /// Reads `bits` bits from the stream and returns a specified unsigned int type
49 /// Reads `bits` bits from the reader and returns a specified type
3250 /// containing them in the least significant end, returning an error if the
3351 /// specified number of bits could not be read.
34 pub fn readBitsNoEof(self: *Self, comptime U: type, bits: usize) !U {
35 var n: usize = undefined;
36 const result = try self.readBits(U, bits, &n);
37 if (n < bits) return error.EndOfStream;
38 return result;
52 pub fn readBitsNoEof(self: *@This(), comptime T: type, num: u16) !T {
53 const b, const c = try self.readBitsTuple(T, num);
54 if (c < num) return error.EndOfStream;
55 return b;
3956 }
4057
41 /// Reads `bits` bits from the stream and returns a specified unsigned int type
58 /// Reads `bits` bits from the reader and returns a specified type
4259 /// containing them in the least significant end. The number of bits successfully
4360 /// 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 {
45 //by extending the buffer to a minimum of u8 we can cover a number of edge cases
46 // related to shifting and casting.
47 const u_bit_count = @bitSizeOf(U);
48 const buf_bit_count = bc: {
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);
61 pub fn readBits(self: *@This(), comptime T: type, num: u16, out_bits: *u16) !T {
62 const b, const c = try self.readBitsTuple(T, num);
63 out_bits.* = c;
64 return b;
65 }
5466
55 out_bits.* = @as(usize, 0);
56 if (U == u0 or bits == 0) return 0;
57 var out_buffer = @as(Buf, 0);
67 /// Reads `bits` bits from the reader and returns a tuple of the specified type
68 /// containing them in the least significant end, and the number of bits successfully
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) {
60 const n = if (self.bit_count >= bits) @as(u3, @intCast(bits)) else self.bit_count;
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
74 //dump any bits in our buffer first
75 if (num <= self.count) return initBits(T, self.removeBits(@intCast(num)), num);
8376
84 //copy bytes until we have enough bits, then leave the rest in bit_buffer
85 while (out_bits.* < bits) {
86 const n = bits - out_bits.*;
87 const next_byte = self.forward_reader.readByte() catch |err| switch (err) {
88 error.EndOfStream => return @as(U, @intCast(out_buffer)),
77 var out_count: u16 = self.count;
78 var out: U = self.removeBits(self.count);
79
80 //grab all the full bytes we need and put their
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),
8987 else => |e| return e,
9088 };
9189
9290 switch (endian) {
9391 .big => {
94 if (n >= u8_bit_count) {
95 out_buffer <<= @as(u3, @intCast(u8_bit_count - 1));
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;
92 if (U == u8) out = 0 else out <<= 8; //shifting u8 by 8 is illegal in Zig
93 out |= byte;
10894 },
10995 .little => {
110 if (n >= u8_bit_count) {
111 out_buffer |= @as(Buf, next_byte) << @as(BufShift, @intCast(out_bits.*));
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;
96 const pos = @as(U, byte) << @intCast(out_count);
97 out |= pos;
12298 },
12399 }
100 out_count += 8;
124101 }
125102
126 return @as(U, @intCast(out_buffer));
127 }
103 const bits_left = num - out_count;
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 {
130 self.bit_buffer = 0;
131 self.bit_count = 0;
108 const final_byte = self.reader.readByte() catch |err| switch (err) {
109 error.EndOfStream => return initBits(T, out, out_count),
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);
132128 }
133129
134 pub fn read(self: *Self, buffer: []u8) Error!usize {
135 var out_bits: usize = undefined;
136 var out_bits_total = @as(usize, 0);
137 //@NOTE: I'm not sure this is a good idea, maybe alignToByte should be forced
138 if (self.bit_count > 0) {
139 for (buffer) |*b| {
140 b.* = try self.readBits(u8, u8_bit_count, &out_bits);
141 out_bits_total += out_bits;
142 }
143 const incomplete_byte = @intFromBool(out_bits_total % u8_bit_count > 0);
144 return (out_bits_total / u8_bit_count) + incomplete_byte;
130 //convenience function for removing bits from
131 //the appropriate part of the buffer based on
132 //endianess.
133 fn removeBits(self: *@This(), num: u4) u8 {
134 if (num == 8) {
135 self.count = 0;
136 return self.bits;
137 }
138
139 const keep = self.count - num;
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),
145147 }
146148
147 return self.forward_reader.read(buffer);
149 self.count = keep;
150 return bits;
148151 }
149152
150 pub fn reader(self: *Self) Reader {
151 return .{ .context = self };
153 pub fn alignToByte(self: *@This()) void {
154 self.bits = 0;
155 self.count = 0;
152156 }
153157 };
154158}
155159
156pub fn bitReader(
157 comptime endian: std.builtin.Endian,
158 underlying_stream: anytype,
159) BitReader(endian, @TypeOf(underlying_stream)) {
160 return BitReader(endian, @TypeOf(underlying_stream)).init(underlying_stream);
160pub fn bitReader(comptime endian: std.builtin.Endian, reader: anytype) BitReader(endian, @TypeOf(reader)) {
161 return .{ .reader = reader };
161162}
162163
164///////////////////////////////
165
163166test "api coverage" {
164167 const mem_be = [_]u8{ 0b11001101, 0b00001011 };
165168 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);
168171 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;
173 const expectError = testing.expectError;
175 const expect = std.testing.expect;
176 const expectError = std.testing.expectError;
174177
175178 try expect(1 == try bit_stream_be.readBits(u2, 1, &out_bits));
176179 try expect(out_bits == 1);
......@@ -186,12 +189,12 @@ test "api coverage" {
186189 try expect(out_bits == 1);
187190
188191 mem_in_be.pos = 0;
189 bit_stream_be.bit_count = 0;
192 bit_stream_be.count = 0;
190193 try expect(0b110011010000101 == try bit_stream_be.readBits(u15, 15, &out_bits));
191194 try expect(out_bits == 15);
192195
193196 mem_in_be.pos = 0;
194 bit_stream_be.bit_count = 0;
197 bit_stream_be.count = 0;
195198 try expect(0b1100110100001011 == try bit_stream_be.readBits(u16, 16, &out_bits));
196199 try expect(out_bits == 16);
197200
......@@ -201,7 +204,7 @@ test "api coverage" {
201204 try expect(out_bits == 0);
202205 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);
205208 var bit_stream_le = bitReader(.little, mem_in_le.reader());
206209
207210 try expect(1 == try bit_stream_le.readBits(u2, 1, &out_bits));
......@@ -218,12 +221,12 @@ test "api coverage" {
218221 try expect(out_bits == 1);
219222
220223 mem_in_le.pos = 0;
221 bit_stream_le.bit_count = 0;
224 bit_stream_le.count = 0;
222225 try expect(0b001010100011101 == try bit_stream_le.readBits(u15, 15, &out_bits));
223226 try expect(out_bits == 15);
224227
225228 mem_in_le.pos = 0;
226 bit_stream_le.bit_count = 0;
229 bit_stream_le.count = 0;
227230 try expect(0b1001010100011101 == try bit_stream_le.readBits(u16, 16, &out_bits));
228231 try expect(out_bits == 16);
229232
lib/std/io/bit_writer.zig+99-114
......@@ -1,153 +1,138 @@
11const 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 stream
8pub fn BitWriter(comptime endian: std.builtin.Endian, comptime WriterType: type) type {
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
16pub fn BitWriter(comptime endian: std.builtin.Endian, comptime Writer: type) type {
917 return struct {
10 forward_writer: WriterType,
11 bit_buffer: u8,
12 bit_count: u4,
13
14 pub const Error = WriterType.Error;
15 pub const Writer = io.Writer(*Self, Error, write);
16
17 const Self = @This();
18 const u8_bit_count = @bitSizeOf(u8);
19 const u4_bit_count = @bitSizeOf(u4);
20
21 pub fn init(forward_writer: WriterType) Self {
22 return Self{
23 .forward_writer = forward_writer,
24 .bit_buffer = 0,
25 .bit_count = 0,
26 };
27 }
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
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
3136 /// are enough to fill a byte.
32 pub fn writeBits(self: *Self, value: anytype, bits: usize) Error!void {
33 if (bits == 0) return;
34
35 const U = @TypeOf(value);
36 comptime assert(@typeInfo(U).int.signedness == .unsigned);
37
38 //by extending the buffer to a minimum of u8 we can cover a number of edge cases
39 // related to shifting and casting.
40 const u_bit_count = @bitSizeOf(U);
41 const buf_bit_count = bc: {
42 assert(u_bit_count >= bits);
43 break :bc if (u_bit_count <= u8_bit_count) u8_bit_count else u_bit_count;
44 };
45 const Buf = std.meta.Int(.unsigned, buf_bit_count);
46 const BufShift = math.Log2Int(Buf);
47
48 const buf_value = @as(Buf, @intCast(value));
49
50 const high_byte_shift = @as(BufShift, @intCast(buf_bit_count - u8_bit_count));
51 var in_buffer = switch (endian) {
52 .big => buf_value << @as(BufShift, @intCast(buf_bit_count - bits)),
53 .little => buf_value,
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));
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
6059 switch (endian) {
6160 .big => {
62 const shift = @as(BufShift, @intCast(high_byte_shift + self.bit_count));
63 const v = @as(u8, @intCast(in_buffer >> shift));
64 self.bit_buffer |= v;
65 in_buffer <<= n;
61 const bits = in >> @intCast(in_count - bits_free);
62 self.addBits(@truncate(bits), bits_free);
6663 },
6764 .little => {
68 const v = @as(u8, @truncate(in_buffer)) << @as(u3, @intCast(self.bit_count));
69 self.bit_buffer |= v;
70 in_buffer >>= n;
65 self.addBits(@truncate(in), bits_free);
66 in >>= @intCast(bits_free);
7167 },
7268 }
73 self.bit_count += n;
74 in_bits -= n;
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;
69 in_count -= bits_free;
70 try self.flushBits();
8171 }
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_buffer
85 while (in_bits >= u8_bit_count) {
73 //write full bytes while we can
74 const full_bytes_left = in_count / 8;
75 for (0..full_bytes_left) |_| {
8676 switch (endian) {
8777 .big => {
88 const v = @as(u8, @intCast(in_buffer >> high_byte_shift));
89 try self.forward_writer.writeByte(v);
90 in_buffer <<= @as(u3, @intCast(u8_bit_count - 1));
91 in_buffer <<= 1;
78 const bits = in >> @intCast(in_count - 8);
79 try self.writer.writeByte(@truncate(bits));
9280 },
9381 .little => {
94 const v = @as(u8, @truncate(in_buffer));
95 try self.forward_writer.writeByte(v);
96 in_buffer >>= @as(u3, @intCast(u8_bit_count - 1));
97 in_buffer >>= 1;
82 try self.writer.writeByte(@truncate(in));
83 if (U == u8) in = 0 else in >>= 8;
9884 },
9985 }
100 in_bits -= u8_bit_count;
86 in_count -= 8;
10187 }
10288
103 if (in_bits > 0) {
104 self.bit_count = @as(u4, @intCast(in_bits));
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;
89 //save the remaining bits in the buffer
90 self.addBits(@truncate(in), @intCast(in_count));
11891 }
11992
120 pub fn write(self: *Self, buffer: []const u8) Error!usize {
121 // TODO: I'm not sure this is a good idea, maybe flushBits should be forced
122 if (self.bit_count > 0) {
123 for (buffer) |b|
124 try self.writeBits(b, u8_bit_count);
125 return buffer.len;
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 },
126105 }
127
128 return self.forward_writer.write(buffer);
106 self.count += num;
129107 }
130108
131 pub fn writer(self: *Self) Writer {
132 return .{ .context = self };
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;
133117 }
134118 };
135119}
136120
137pub fn bitWriter(
138 comptime endian: std.builtin.Endian,
139 underlying_stream: anytype,
140) BitWriter(endian, @TypeOf(underlying_stream)) {
141 return BitWriter(endian, @TypeOf(underlying_stream)).init(underlying_stream);
121pub fn bitWriter(comptime endian: std.builtin.Endian, writer: anytype) BitWriter(endian, @TypeOf(writer)) {
122 return .{ .writer = writer };
142123}
143124
125///////////////////////////////
126
144127test "api coverage" {
145128 var mem_be = [_]u8{0} ** 2;
146129 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);
149132 var bit_stream_be = bitWriter(.big, mem_out_be.writer());
150133
134 const testing = std.testing;
135
151136 try bit_stream_be.writeBits(@as(u2, 1), 1);
152137 try bit_stream_be.writeBits(@as(u5, 2), 2);
153138 try bit_stream_be.writeBits(@as(u128, 3), 3);
......@@ -169,7 +154,7 @@ test "api coverage" {
169154
170155 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);
173158 var bit_stream_le = bitWriter(.little, mem_out_le.writer());
174159
175160 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" {
8282
8383 var bit_stream = io.bitReader(native_endian, file.reader());
8484
85 var out_bits: usize = undefined;
85 var out_bits: u16 = undefined;
8686
8787 try expect(1 == try bit_stream.readBits(u2, 1, &out_bits));
8888 try expect(out_bits == 1);