| ... | @@ -1,238 +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 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 |
| 16 | pub fn BitReader(comptime endian: std.builtin.Endian, comptime Reader: type) type { |
| 17 | return struct { |
| 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, |
| 46 | }; |
| 47 | } |
| 48 | |
| 49 | /// Reads `bits` bits from the reader and returns a specified type |
| 50 | /// containing them in the least significant end, returning an error if the |
| 51 | /// specified number of bits could not be read. |
| 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; |
| 56 | } |
| 57 | |
| 58 | /// Reads `bits` bits from the reader and returns a specified type |
| 59 | /// containing them in the least significant end. The number of bits successfully |
| 60 | /// read is placed in `out_bits`, as reaching the end of the stream is not an error. |
| 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 | } |
| 66 | |
| 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 |
| 73 | |
| 74 | //dump any bits in our buffer first |
| 75 | if (num <= self.count) return initBits(T, self.removeBits(@intCast(num)), num); |
| 76 | |
| 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), |
| 87 | else => |e| return e, |
| 88 | }; |
| 89 | |
| 90 | switch (endian) { |
| 91 | .big => { |
| 92 | if (U == u8) out = 0 else out <<= 8; //shifting u8 by 8 is illegal in Zig |
| 93 | out |= byte; |
| 94 | }, |
| 95 | .little => { |
| 96 | const pos = @as(U, byte) << @intCast(out_count); |
| 97 | out |= pos; |
| 98 | }, |
| 99 | } |
| 100 | out_count += 8; |
| 101 | } |
| 102 | |
| 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); |
| 107 | |
| 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); |
| 128 | } |
| 129 | |
| 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), |
| 147 | } |
| 148 | |
| 149 | self.count = keep; |
| 150 | return bits; |
| 151 | } |
| 152 | |
| 153 | pub fn alignToByte(self: *@This()) void { |
| 154 | self.bits = 0; |
| 155 | self.count = 0; |
| 156 | } |
| 157 | }; |
| 158 | } |
| 159 | |
| 160 | pub fn bitReader(comptime endian: std.builtin.Endian, reader: anytype) BitReader(endian, @TypeOf(reader)) { |
| 161 | return .{ .reader = reader }; |
| 162 | } |
| 163 | |
| 164 | /////////////////////////////// |
| 165 | |
| 166 | test "api coverage" { |
| 167 | const mem_be = [_]u8{ 0b11001101, 0b00001011 }; |
| 168 | const mem_le = [_]u8{ 0b00011101, 0b10010101 }; |
| 169 | |
| 170 | var mem_in_be = std.io.fixedBufferStream(&mem_be); |
| 171 | var bit_stream_be = bitReader(.big, mem_in_be.reader()); |
| 172 | |
| 173 | var out_bits: u16 = undefined; |
| 174 | |
| 175 | const expect = std.testing.expect; |
| 176 | const expectError = std.testing.expectError; |
| 177 | |
| 178 | try expect(1 == try bit_stream_be.readBits(u2, 1, &out_bits)); |
| 179 | try expect(out_bits == 1); |
| 180 | try expect(2 == try bit_stream_be.readBits(u5, 2, &out_bits)); |
| 181 | try expect(out_bits == 2); |
| 182 | try expect(3 == try bit_stream_be.readBits(u128, 3, &out_bits)); |
| 183 | try expect(out_bits == 3); |
| 184 | try expect(4 == try bit_stream_be.readBits(u8, 4, &out_bits)); |
| 185 | try expect(out_bits == 4); |
| 186 | try expect(5 == try bit_stream_be.readBits(u9, 5, &out_bits)); |
| 187 | try expect(out_bits == 5); |
| 188 | try expect(1 == try bit_stream_be.readBits(u1, 1, &out_bits)); |
| 189 | try expect(out_bits == 1); |
| 190 | |
| 191 | mem_in_be.pos = 0; |
| 192 | bit_stream_be.count = 0; |
| 193 | try expect(0b110011010000101 == try bit_stream_be.readBits(u15, 15, &out_bits)); |
| 194 | try expect(out_bits == 15); |
| 195 | |
| 196 | mem_in_be.pos = 0; |
| 197 | bit_stream_be.count = 0; |
| 198 | try expect(0b1100110100001011 == try bit_stream_be.readBits(u16, 16, &out_bits)); |
| 199 | try expect(out_bits == 16); |
| 200 | |
| 201 | _ = try bit_stream_be.readBits(u0, 0, &out_bits); |
| 202 | |
| 203 | try expect(0 == try bit_stream_be.readBits(u1, 1, &out_bits)); |
| 204 | try expect(out_bits == 0); |
| 205 | try expectError(error.EndOfStream, bit_stream_be.readBitsNoEof(u1, 1)); |
| 206 | |
| 207 | var mem_in_le = std.io.fixedBufferStream(&mem_le); |
| 208 | var bit_stream_le = bitReader(.little, mem_in_le.reader()); |
| 209 | |
| 210 | try expect(1 == try bit_stream_le.readBits(u2, 1, &out_bits)); |
| 211 | try expect(out_bits == 1); |
| 212 | try expect(2 == try bit_stream_le.readBits(u5, 2, &out_bits)); |
| 213 | try expect(out_bits == 2); |
| 214 | try expect(3 == try bit_stream_le.readBits(u128, 3, &out_bits)); |
| 215 | try expect(out_bits == 3); |
| 216 | try expect(4 == try bit_stream_le.readBits(u8, 4, &out_bits)); |
| 217 | try expect(out_bits == 4); |
| 218 | try expect(5 == try bit_stream_le.readBits(u9, 5, &out_bits)); |
| 219 | try expect(out_bits == 5); |
| 220 | try expect(1 == try bit_stream_le.readBits(u1, 1, &out_bits)); |
| 221 | try expect(out_bits == 1); |
| 222 | |
| 223 | mem_in_le.pos = 0; |
| 224 | bit_stream_le.count = 0; |
| 225 | try expect(0b001010100011101 == try bit_stream_le.readBits(u15, 15, &out_bits)); |
| 226 | try expect(out_bits == 15); |
| 227 | |
| 228 | mem_in_le.pos = 0; |
| 229 | bit_stream_le.count = 0; |
| 230 | try expect(0b1001010100011101 == try bit_stream_le.readBits(u16, 16, &out_bits)); |
| 231 | try expect(out_bits == 16); |
| 232 | |
| 233 | _ = try bit_stream_le.readBits(u0, 0, &out_bits); |
| 234 | |
| 235 | try expect(0 == try bit_stream_le.readBits(u1, 1, &out_bits)); |
| 236 | try expect(out_bits == 0); |
| 237 | try expectError(error.EndOfStream, bit_stream_le.readBitsNoEof(u1, 1)); |
| 238 | } |