From eeb67c3b79a6b9a8a02582642bfa3bf1671db49f Mon Sep 17 00:00:00 2001
From: Muaz Ahmad <mahmad2000@protonmail.com>
Date: Thu, 7 Dec 2023 14:33:27 +0500
Subject: [PATCH] scan writing done

---
 src/output.zig | 76 +++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 66 insertions(+), 10 deletions(-)

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;