From 6c91c2c228f2e8b7083079791c36c92d3da709a1 Mon Sep 17 00:00:00 2001
From: Muaz Ahmad <mahmad2000@protonmail.com>
Date: Tue, 12 Dec 2023 12:19:53 +0500
Subject: [PATCH] bitwriter reimplemented

---
 src/output.zig | 10 +++++-----
 src/util.zig   | 46 ++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/src/output.zig b/src/output.zig
index f985273..c1dac8e 100644
--- a/src/output.zig
+++ b/src/output.zig
@@ -48,14 +48,14 @@ const RLEWriter = struct {
     fn write_value(self: *Self, huffcode: HuffCode, unit: RLE_Unit) !void {
         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));
+        const value_size = unit.symbol & 0x0f;
+        if (value_size != 0) {
+            try self.bw.write_bits(unit_val & (try std.math.powi(u16, 2, value_size) - 1), @truncate(value_size));
+        }
     }
 
     inline fn flush(self: *Self) !void {
-        if (self.bw.bits_used != 0) {
-            const bits_left = 8 - self.bw.bits_used;
-            try self.bw.write_bits(0xf, bits_left);
-        }
+        try self.bw.flush_end();
     }
 };
 
diff --git a/src/util.zig b/src/util.zig
index d3eaaff..9ed819f 100644
--- a/src/util.zig
+++ b/src/util.zig
@@ -149,7 +149,7 @@ pub fn gen_qtable(q: f16, step_start_band: usize, step_stop_band: usize) @Vector
 pub const BufferedBitWriter = struct {
     byte_buff: u8,
     bits_used: u4,
-    buffer: [32]u8,
+    buffer: [8]u8,
     buffer_idx: usize,
     f: std.fs.File,
 
@@ -160,18 +160,52 @@ pub const BufferedBitWriter = struct {
             .f = f,
             .byte_buff = 0,
             .bits_used = 0,
-            .buffer = [_]u8{0x00} ** 32,
+            .buffer = [_]u8{0x00} ** 8,
             .buffer_idx = 0,
         };
     }
 
     pub fn write_bits(self: *Self, val: u16, n_bits: u5) !void {
-        _ = n_bits;
-        _ = val;
-        _ = self;
+        const curr_byte_space = 8 - self.bits_used;
+        if (n_bits <= curr_byte_space) {
+            try self.add_bits(val, @truncate(n_bits));
+        } else {
+            try self.add_bits(val >> @truncate(n_bits - curr_byte_space), @truncate(curr_byte_space));
+            const val_remaining_mask = try std.math.powi(u16, 2, n_bits - curr_byte_space) - 1;
+            try self.write_bits(val & val_remaining_mask, n_bits - curr_byte_space);
+        }
+    }
+
+    inline fn add_bits(self: *Self, val: u16, n_bits: u4) !void {
+        self.byte_buff |= @truncate(val << (8 - self.bits_used - n_bits));
+        self.bits_used += n_bits;
+        if (self.bits_used == 8) {
+            self.buffer[self.buffer_idx] = self.byte_buff;
+            if (self.byte_buff == 0xff) {
+                self.buffer_idx += 1;
+            }
+            self.buffer_idx += 1;
+            self.bits_used = 0;
+            self.byte_buff = 0x00;
+            if (self.buffer_idx >= self.buffer.len) {
+                try self.flush();
+                self.buffer_idx %= self.buffer.len;
+            }
+        }
     }
 
     pub fn flush(self: *Self) !void {
-        _ = self;
+        _ = try self.f.write(&self.buffer);
+        @memset(&self.buffer, 0);
+    }
+
+    pub fn flush_end(self: *Self) !void {
+        if (self.bits_used != 0) {
+            const stuffing = try std.math.powi(u8, 2, 8 - self.bits_used) - 1;
+            try self.add_bits(stuffing, 8 - self.bits_used);
+        }
+        if (self.buffer_idx != 0) {
+            _ = try self.f.write(self.buffer[0..self.buffer_idx]);
+        }
     }
 };