diff --git a/src/output.zig b/src/output.zig index 2408600..f985273 100644 --- a/src/output.zig +++ b/src/output.zig @@ -19,20 +19,14 @@ const HuffmanMeta = struct { huffval: [256]u8, }; -const BitWriter = struct { - f: std.fs.File, - buff: [8]u8, - curr_buff_idx: usize, - curr_bit_used: u4, +const RLEWriter = struct { + bw: util.BufferedBitWriter, const Self = @This(); fn init(f: std.fs.File) Self { return Self{ - .f = f, - .buff = [_]u8{0x00} ** 8, - .curr_buff_idx = 0, - .curr_bit_used = 0, + .bw = util.BufferedBitWriter.init(f), }; } @@ -42,75 +36,27 @@ const BitWriter = struct { try self.write_value(dc_huff.get(units[i].symbol).?, units[i]); i += 1; - while (units[i].symbol != 0x00) : (i += 1) { + while (true) : (i += 1) { try self.write_value(ac_huff.get(units[i].symbol).?, units[i]); + if (units[i].symbol == 0x00) { + i += 1; + break; + } } - try self.write_value(ac_huff.get(units[i].symbol).?, units[i]); - i += 1; } 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); - } + try self.bw.write_bits(huffcode.value, huffcode.n_bits); + const unit_val: u16 = if (unit.value >= 0) @bitCast(unit.value) else @bitCast(unit.value - 1); + try self.bw.write_bits(unit_val, @truncate(unit.symbol & 0x0f)); } 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); + if (self.bw.bits_used != 0) { + const bits_left = 8 - self.bw.bits_used; + try self.bw.write_bits(0xf, bits_left); } } - - 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]); - } }; const Scan = struct { @@ -152,7 +98,6 @@ const Scan = struct { for (0..4) |i| { @memset(self.freqs[i], 0); } - for (0..h) |x| { for (0..w) |y| { for (0..4) |Y_sub_i| { @@ -175,7 +120,7 @@ const Scan = struct { fn dump_scan(self: *Self, f: std.fs.File) !void { _ = try f.write(&[_]u8{ 0xff, 0xda, 0x00, 0x0c, 0x03, 0x01, 0x00, 0x02, 0x11, 0x03, 0x11, 0x00, 0x3f, 0x00 }); - var bw = BitWriter.init(f); + var bw = RLEWriter.init(f); var idxs = [3]usize{ 0, 0, 0 }; while (idxs[0] < self.rles[0].items.len and idxs[1] < self.rles[1].items.len and idxs[2] < self.rles[2].items.len) { for (0..4) |_| { @@ -184,7 +129,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 bw.flush(); _ = try f.write(&[2]u8{ 0xff, 0xd9 }); } }; @@ -392,6 +337,7 @@ fn parse_block(block: *util.BlockQuantized, dc_diff: *i16, rle: *RLE_Seq, dc_fre .symbol = symbol, .value = block[i], }); + curr_rlen = 0; } ac_freqs[0x00] += 1; try rle.append(RLE_Unit{ diff --git a/src/util.zig b/src/util.zig index 78aa388..d3eaaff 100644 --- a/src/util.zig +++ b/src/util.zig @@ -112,8 +112,8 @@ pub const Buffers = struct { buffs.U_quant[i] = try alloc.alloc(BlockQuantized, block_w); buffs.V_quant[i] = try alloc.alloc(BlockQuantized, block_w); } - buffs.Q_Lum.* = gen_qtable(255, q, 12); - buffs.Q_Chrom.* = gen_qtable(255, q, 6); + buffs.Q_Lum.* = gen_qtable(q, 5, 12); + buffs.Q_Chrom.* = gen_qtable(q, 5, 10); return buffs; } @@ -122,15 +122,20 @@ pub const Buffers = struct { } }; -pub fn gen_qtable(q_max: usize, q: f16, band_range: usize) @Vector(64, f16) { - var ret: @Vector(64, f16) = [_]f16{0.0} ** 64; - const q_min: usize = @intFromFloat(255 - 250 * q); +pub fn gen_qtable(q: f16, step_start_band: usize, step_stop_band: usize) @Vector(64, f16) { + var ret: @Vector(64, f16) = [_]f32{0.0} ** 64; + const band_range = step_stop_band - step_start_band; + const q_max: usize = @intFromFloat(255 - 235 * q); + const q_min = 8; var step = (q_max - q_min) / band_range; var idx: usize = 0; for (0..16) |band_i| { const band_len = if (band_i < 8) band_i + 1 else 15 - band_i; for (0..band_len) |j| { - if (band_i < band_range) { + if (band_i < step_start_band) { + ret[idx + j] = @floatFromInt(q_min); + } + if (band_i < band_range + step_start_band) { ret[idx + j] = @floatFromInt(q_min + step * band_i); } else { ret[idx + j] = @floatFromInt(q_max); @@ -140,3 +145,33 @@ pub fn gen_qtable(q_max: usize, q: f16, band_range: usize) @Vector(64, f16) { } return ret; } + +pub const BufferedBitWriter = struct { + byte_buff: u8, + bits_used: u4, + buffer: [32]u8, + buffer_idx: usize, + f: std.fs.File, + + const Self = @This(); + + pub fn init(f: std.fs.File) Self { + return Self{ + .f = f, + .byte_buff = 0, + .bits_used = 0, + .buffer = [_]u8{0x00} ** 32, + .buffer_idx = 0, + }; + } + + pub fn write_bits(self: *Self, val: u16, n_bits: u5) !void { + _ = n_bits; + _ = val; + _ = self; + } + + pub fn flush(self: *Self) !void { + _ = self; + } +};