From af6d436079de8f28f17efac1c783c67db7111355 Mon Sep 17 00:00:00 2001
From: Muaz Ahmad <mahmad2000@protonmail.com>
Date: Tue, 28 Nov 2023 13:44:46 +0500
Subject: [PATCH] restructured to allow qtables

---
 src/main.zig    |  4 ++--
 src/threads.zig |  8 ++++----
 src/util.zig    | 31 ++++++++++++++++++++++++++++++-
 3 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/src/main.zig b/src/main.zig
index 009f410..1555867 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -58,10 +58,10 @@ fn get_opts() !util.Options {
 pub fn main() !void {
     const opts = try get_opts();
 
-    var buffs = try util.Buffers.init(std.heap.page_allocator, opts.width, opts.height, 1000);
+    var buffs = try util.Buffers.init(std.heap.page_allocator, opts.width, opts.height, 1000, opts.quality);
     defer buffs.deinit();
 
-    var thread_manager = try threads.ThreadManager.init(std.heap.page_allocator, opts.n_quant_jobs);
+    var thread_manager = try threads.ThreadManager.init(std.heap.page_allocator, opts.n_quant_jobs, buffs.Q_Lum, buffs.Q_Chrom);
     defer thread_manager.deinit();
 
     var f = std.io.getStdIn();
diff --git a/src/threads.zig b/src/threads.zig
index b3490d1..527df87 100644
--- a/src/threads.zig
+++ b/src/threads.zig
@@ -45,7 +45,7 @@ pub const ThreadManager = struct {
 
     const Self = @This();
 
-    pub fn init(alloc: std.mem.Allocator, n_quant_jobs: usize) !Self {
+    pub fn init(alloc: std.mem.Allocator, n_quant_jobs: usize, Q_Lum: *util.QTable, Q_Chrom: *util.QTable) !Self {
         var arena = std.heap.ArenaAllocator.init(alloc);
         var arena_alloc = arena.allocator();
         var job_pool = util.JobPool.init(alloc);
@@ -59,7 +59,7 @@ pub const ThreadManager = struct {
         thread_mgr.signals.* = Signals.init();
 
         for (0..n_quant_jobs) |_| {
-            try thread_mgr.threads.append(try std.Thread.spawn(.{}, quantize_loop, .{ thread_mgr.queue_wrp.queue, thread_mgr.signals }));
+            try thread_mgr.threads.append(try std.Thread.spawn(.{}, quantize_loop, .{ thread_mgr.queue_wrp.queue, thread_mgr.signals, Q_Lum, Q_Chrom }));
         }
 
         return thread_mgr;
@@ -71,10 +71,10 @@ pub const ThreadManager = struct {
     }
 };
 
-fn quantize_loop(queue: *util.JobQueue, signals: *Signals) void {
+fn quantize_loop(queue: *util.JobQueue, signals: *Signals, Q_Lum: *util.QTable, Q_Chrom: *util.QTable) void {
     while (queue.HasJobs() or !signals.quit.load(std.builtin.AtomicOrder.Acquire)) : (std.time.sleep(1)) {
         const job = queue.pop() orelse continue;
-        transform.quantize(job.source, job.target);
+        transform.quantize(job.source, job.target, if (job.is_lum) Q_Lum else Q_Chrom);
         _ = @atomicRmw(u32, &signals.processed.value, .Add, 1, .SeqCst);
     }
 }
diff --git a/src/util.zig b/src/util.zig
index 928a2dc..340e672 100644
--- a/src/util.zig
+++ b/src/util.zig
@@ -9,11 +9,13 @@ pub const Options = struct {
 
 pub const Block = [64]u8;
 pub const BlockQuantized = [64]i16;
+pub const QTable = @Vector(64, f16);
 
 pub const ThreadList = std.ArrayList(std.Thread);
 pub const Job = struct {
     source: *Block,
     target: *BlockQuantized,
+    is_lum: bool,
 };
 pub const JobQueue = struct {
     const List = std.TailQueue(Job);
@@ -65,6 +67,9 @@ pub const Buffers = struct {
     U: [][]Block,
     V: [][]Block,
 
+    Q_Lum: *QTable,
+    Q_Chrom: *QTable,
+
     Y_quant: [][][4]BlockQuantized,
     U_quant: [][]BlockQuantized,
     V_quant: [][]BlockQuantized,
@@ -76,7 +81,7 @@ pub const Buffers = struct {
 
     const Self = @This();
 
-    pub fn init(root_alloc: std.mem.Allocator, w: usize, h: usize, output_buff_len: usize) !Self {
+    pub fn init(root_alloc: std.mem.Allocator, w: usize, h: usize, output_buff_len: usize, q: f16) !Self {
         var arena = std.heap.ArenaAllocator.init(root_alloc);
         var alloc = arena.allocator();
         const block_w = w / 16;
@@ -88,6 +93,9 @@ pub const Buffers = struct {
             .U = try alloc.alloc([]Block, block_h),
             .V = try alloc.alloc([]Block, block_h),
 
+            .Q_Lum = try alloc.create(QTable),
+            .Q_Chrom = try alloc.create(QTable),
+
             .Y_quant = try alloc.alloc([][4]BlockQuantized, block_h),
             .U_quant = try alloc.alloc([]BlockQuantized, block_h),
             .V_quant = try alloc.alloc([]BlockQuantized, block_h),
@@ -106,6 +114,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);
         return buffs;
     }
 
@@ -113,3 +123,22 @@ pub const Buffers = struct {
         self.arena.deinit();
     }
 };
+
+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 - 100 * q);
+    var step = (q_max - q_min) / band_range;
+    var idx: usize = 0;
+    for (0..16) |band_i| {
+        const band_len = band_i + 1;
+        for (0..band_len) |j| {
+            if (band_i < band_range) {
+                ret[idx + j] = @floatFromInt(q_min + step * band_i);
+            } else {
+                ret[idx + j] = @floatFromInt(q_max);
+            }
+        }
+        idx += band_len;
+    }
+    return ret;
+}