From 125a3106d6f0e21e3458b138fd9cce818928e336 Mon Sep 17 00:00:00 2001
From: siivonek <kari.siivonen@tuni.fi>
Date: Wed, 24 Aug 2022 18:45:13 +0300
Subject: [PATCH] [isp] Use temporary coeff array when quantizing coeffs. After
 deriving coeffs, copy temp coeffs from linear order to correct arrays with cu
 order.

---
 src/transform.c | 56 ++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 42 insertions(+), 14 deletions(-)

diff --git a/src/transform.c b/src/transform.c
index b9fabd65..c5a38475 100644
--- a/src/transform.c
+++ b/src/transform.c
@@ -1107,6 +1107,19 @@ int uvg_quantize_residual_trskip(
   return best->has_coeffs;
 }
 
+
+static INLINE int translate_to_cu_order_idx(const int lcu_x, const int lcu_y, const int block_w, const int block_h, const int linear_idx)
+{
+  // ISP_TODO: get rid of all there temp variables after making sure this works
+  const int start_idx = lcu_x + lcu_y * LCU_WIDTH;
+  const int offset_x = linear_idx % block_w;
+  const int local_y = linear_idx / block_h;
+  const int offset_y = local_y * LCU_WIDTH;
+
+  return (start_idx + offset_x + offset_y);
+}
+
+
 /**
  * Calculate the residual coefficients for a single TU.
  *
@@ -1124,6 +1137,7 @@ static void quantize_tr_residual(
 {
   const int x = cu_loc->x;
   const int y = cu_loc->y;
+
   const uvg_config *cfg    = &state->encoder_control->cfg;
   const int32_t shift      = color == COLOR_Y ? 0 : 1;
   const vector2d_t lcu_px  = { SUB_SCU(x) >> shift, SUB_SCU(y) >> shift};
@@ -1158,26 +1172,27 @@ static void quantize_tr_residual(
   uvg_pixel *pred = NULL;
   // Pointers to current location in arrays with reference.
   const uvg_pixel *ref = NULL;
-  // Pointers to current location in arrays with quantized coefficients.
-  coeff_t *coeff = NULL;
+  // Temp coeff array
+  coeff_t coeff[TR_MAX_WIDTH * TR_MAX_WIDTH];
+  coeff_t *dst_coeff = NULL;
 
   // ISP_TODO: use temp coeff array size MAX_TR_WIDTH^2 instead of coeff pointers
-  // ISP_TODO: inside temp coeff array, entries are in the old order. PÖTKÖ
+  // ISP_TODO: inside temp coeff array, entries are in the old linear order. PÖTKÖ
   switch (color) {
     case COLOR_Y:
-      pred  = &lcu->rec.y[offset];
-      ref   = &lcu->ref.y[offset];
-      coeff = &lcu->coeff.y[z_index];
+      pred      = &lcu->rec.y[offset];
+      ref       = &lcu->ref.y[offset];
+      dst_coeff = &lcu->coeff.y;
       break;
     case COLOR_U:
-      pred = &lcu->rec.u[offset];
-      ref  = &lcu->ref.u[offset];
-      coeff = &lcu->coeff.u[z_index];
+      pred      = &lcu->rec.u[offset];
+      ref       = &lcu->ref.u[offset];
+      dst_coeff = &lcu->coeff.u;
       break;
     case COLOR_V:
-      pred = &lcu->rec.v[offset];
-      ref  = &lcu->ref.v[offset];
-      coeff = &lcu->coeff.v[z_index];
+      pred      = &lcu->rec.v[offset];
+      ref       = &lcu->ref.v[offset];
+      dst_coeff = &lcu->coeff.v;
       break;
     default:
       break;
@@ -1274,9 +1289,22 @@ static void quantize_tr_residual(
 
   cbf_clear(&cur_pu->cbf, depth, color);
   if (has_coeffs) {
-    // ISP_TODO: copy coeffs into CU order instead of pötkö
+    const int coeffs_to_copy = tr_width * tr_height;
+    for (int i = 0; i < coeffs_to_copy; ++i) {
+      const coeff_t c = coeff[i];
+      const idx = translate_to_cu_order_idx(lcu_px.x, lcu_px.y, tr_width, tr_height, i);
+      dst_coeff[idx] = c;
+    }
     cbf_set(&cur_pu->cbf, depth, color);
-  } // ISP_TODO: if no coeffs, mem set width * height amount of coeffs to zero
+  }
+  else {
+    // ISP_TODO: if no coeffs, mem set width * height amount of coeffs to zero
+    int idx = lcu_px.x + lcu_px.y * LCU_WIDTH;
+    for (int j = 0; j < tr_height; ++j) {
+      memset(dst_coeff[idx], 0, (sizeof(coeff_t) * tr_width));
+      idx += LCU_WIDTH;
+    }
+  }
 }
 
 /**