diff --git a/src/output.zig b/src/output.zig index 54443da..2408600 100644 --- a/src/output.zig +++ b/src/output.zig @@ -40,20 +40,76 @@ const BitWriter = struct { var i = idx_ptr.*; defer idx_ptr.* = i; - try self.write_value(dc_huff.get(units[i].symbol).?, if (units[i].value >= 0) units[i].value else -units[i].value, units[i].symbol & 0x0f); + try self.write_value(dc_huff.get(units[i].symbol).?, units[i]); i += 1; while (units[i].symbol != 0x00) : (i += 1) { - try self.write_value(ac_huff.get(units[i].symbol).?, units[i].value, units[i].symbol & 0x0f); + try self.write_value(ac_huff.get(units[i].symbol).?, units[i]); } - try self.write_value(ac_huff.get(units[i].symbol).?, units[i].value, units[i].symbol & 0x0f); + try self.write_value(ac_huff.get(units[i].symbol).?, units[i]); i += 1; } - fn write_value(self: *Self, huffcode: HuffCode, value: i16, val_bits: u8) !void { - _ = val_bits; - _ = value; - _ = huffcode; - _ = self; + fn write_value(self: *Self, huffcode: HuffCode, unit: RLE_Unit) !void { + try self.write_bits(huffcode.value, huffcode.n_bits); + const unit_val: u16 = if (unit.value >= 0) blk: { + break :blk @bitCast(unit.value); + } else blk: { + const tmp: u16 = @bitCast(unit.value - 1); + break :blk tmp & (try std.math.powi(u16, 2, unit.symbol & 0x0f) - 1); + }; + try self.write_bits(unit_val, unit.symbol & 0x0f); + } + + inline fn write_bits(self: *Self, val: u16, n_bits: u8) !void { + if (self.curr_bit_used + n_bits < 8) { + self.buff[self.curr_buff_idx] = @truncate((self.buff[self.curr_buff_idx] << @truncate(n_bits)) | val); + self.curr_bit_used += @truncate(n_bits); + } else { + const bits_align_n = 8 - self.curr_bit_used; + self.buff[self.curr_buff_idx] = @truncate((self.buff[self.curr_buff_idx] << @truncate(bits_align_n)) | (val >> @truncate(n_bits - bits_align_n))); + var byte_stuff_flag = self.buff[self.curr_buff_idx] == 0xff; + self.curr_bit_used = 0; + self.curr_buff_idx += 1; + try self.flush(); + if (byte_stuff_flag) { + self.curr_buff_idx += 1; + try self.flush(); + } + + var bits_left = n_bits - bits_align_n; + var val_left = val & (try std.math.powi(u16, 2, bits_left) - 1); + + while (bits_left >= 8) : (bits_left -= 8) { + self.buff[self.curr_buff_idx] = @truncate(val_left >> @truncate(bits_left - 8)); + byte_stuff_flag = self.buff[self.curr_buff_idx] == 0xff; + val_left = val_left & (try std.math.powi(u16, 2, bits_left - 8) - 1); + self.curr_buff_idx += 1; + if (self.curr_buff_idx == self.buff.len) try self.flush(); + if (byte_stuff_flag) { + self.curr_buff_idx += 1; + try self.flush(); + } + } + self.buff[self.curr_buff_idx] = @truncate((self.buff[self.curr_buff_idx] << @truncate(bits_left)) | val_left); + self.curr_bit_used += @truncate(bits_left); + } + } + + inline fn flush(self: *Self) !void { + if (self.curr_buff_idx == self.buff.len) { + _ = try self.f.write(&self.buff); + self.curr_buff_idx = 0; + @memset(&self.buff, 0); + } + } + + inline fn flush_end(self: *Self) !void { + if (self.curr_bit_used != 0) { + const stuff_n = 8 - self.curr_bit_used; + self.buff[self.curr_buff_idx] = @truncate((self.curr_buff_idx << stuff_n) | (try std.math.powi(u8, 2, stuff_n) - 1)); + self.curr_buff_idx += 1; + } + _ = try self.f.write(self.buff[0..self.curr_buff_idx]); } }; @@ -128,7 +184,7 @@ const Scan = struct { try bw.write_until_eob(self.rles[1].items, &idxs[1], &self.huffs[2], &self.huffs[3]); try bw.write_until_eob(self.rles[2].items, &idxs[2], &self.huffs[2], &self.huffs[3]); } - + try bw.flush_end(); _ = try f.write(&[2]u8{ 0xff, 0xd9 }); } }; @@ -186,7 +242,7 @@ inline fn gen_codes(huff: *Huffman, huffman_meta: HuffmanMeta) !void { j = 1; } huffsize[k] = 0; - const total_k = k + 1; + const total_k = k; k = 0; var code: usize = 0;