From 491f6a98b7bdf26a970975de3f8c96c2f6673222 Mon Sep 17 00:00:00 2001
From: Muaz Ahmad <mahmad2000@protonmail.com>
Date: Fri, 24 Nov 2023 15:40:22 +0500
Subject: [PATCH] added allocation for permanent buffers

---
 src/main.zig |  8 ++++----
 src/util.zig | 57 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/src/main.zig b/src/main.zig
index 9650d2d..6a72c9b 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -54,10 +54,10 @@ fn get_opts() !util.Options {
 }
 
 pub fn main() !void {
-    const opts = get_opts();
+    const opts = try get_opts();
 
-    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
-    defer arena.deinit();
-    
+    var buffs = try util.Buffers.init(std.heap.page_allocator, opts.width, opts.height, 1000);
+    defer buffs.deinit();
+        
     std.debug.print("{any}\n", .{opts});
 }
diff --git a/src/util.zig b/src/util.zig
index 7b94454..dcc845b 100644
--- a/src/util.zig
+++ b/src/util.zig
@@ -1,6 +1,63 @@
+const std = @import("std");
+
 pub const Options = struct {
     width: usize,
     height: usize,
     quality: f16,
     n_quant_jobs: usize,
 };
+
+pub const Block = [64]u8;
+pub const BlockQuantized = [64]i16;
+
+pub const Buffers = struct {
+    arena: std.heap.ArenaAllocator,
+    
+    Y: [][][4]Block,
+    U: [][]Block,
+    V: [][]Block,
+
+    Y_quant: [][][4]BlockQuantized,
+    U_quant: [][]BlockQuantized,
+    V_quant: [][]BlockQuantized,
+
+    input_buff: []u8,
+    output_buff: []u8,
+
+    const Self = @This();
+
+    pub fn init(root_alloc: std.mem.Allocator, w: usize, h: usize, output_buff_len: usize) !Self {
+        var arena = std.heap.ArenaAllocator.init(root_alloc);
+        var alloc = arena.allocator();
+        const block_w = w / 16;
+        const block_h = h / 16;
+        var buffs = Self {
+            .arena = arena,
+            
+            .Y = try alloc.alloc([][4]Block, block_h),
+            .U = try alloc.alloc([]Block, block_h),
+            .V = try alloc.alloc([]Block, block_h),
+
+            .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),
+
+            .input_buff = try alloc.alloc(u8, w * 8),
+            .output_buff = try alloc.alloc(u8, output_buff_len),
+        };
+        for (0..block_h) |i| {
+            buffs.Y[i] = try alloc.alloc([4]Block, block_w);
+            buffs.U[i] = try alloc.alloc(Block, block_w);
+            buffs.V[i] = try alloc.alloc(Block, block_w);
+
+            buffs.Y_quant[i] = try alloc.alloc([4]BlockQuantized, block_w);
+            buffs.U_quant[i] = try alloc.alloc(BlockQuantized, block_w);
+            buffs.V_quant[i] = try alloc.alloc(BlockQuantized, block_w);
+        }
+        return buffs;
+    }
+
+    pub fn deinit(self: *Self) void {
+        self.arena.deinit();
+    }
+};