From df01d5488f77e9b74ee12c716efe4586cd140d2a Mon Sep 17 00:00:00 2001
From: Joose Sainio <joose.sainio@tuni.fi>
Date: Mon, 14 Nov 2022 10:16:25 +0200
Subject: [PATCH] [quant] fix fast coeff cost

---
 src/rdo.c | 50 +++++++++++++++++++++++++-------------------------
 1 file changed, 25 insertions(+), 25 deletions(-)

diff --git a/src/rdo.c b/src/rdo.c
index fa7506b8..0ed561e1 100644
--- a/src/rdo.c
+++ b/src/rdo.c
@@ -301,35 +301,17 @@ static INLINE double get_coeff_cabac_cost(
   color_t color,
   int8_t scan_mode,
   int8_t tr_skip,
-  cu_info_t* cur_tu,
-  int coeff_order)
+  cu_info_t* cur_tu)
 {
   const int width  = cu_loc->width;
   const int height = cu_loc->height;
   const int sub_coeff_w = color == COLOR_Y ? cu_loc->width  : cu_loc->chroma_width;
   const int sub_coeff_h = color == COLOR_Y ? cu_loc->height : cu_loc->chroma_height;
-  const int lcu_width = color == COLOR_Y ? LCU_WIDTH : LCU_WIDTH_C;
-
-  int x_local = cu_loc->x % LCU_WIDTH;
-  int y_local = cu_loc->y % LCU_WIDTH;
 
   // Make sure there are coeffs present
   bool found = false;
-
-  coeff_t* coeff_ptr = NULL;
-  coeff_t sub_coeff[TR_MAX_WIDTH * TR_MAX_WIDTH];
-
-  if (coeff_order == COEFF_ORDER_LINEAR) {
-    coeff_ptr = (coeff_t*)coeff;
-  }
-  else {
-    // Coeff order CU
-    uvg_get_sub_coeff(sub_coeff, coeff, x_local, y_local, sub_coeff_w, sub_coeff_h, lcu_width);
-    coeff_ptr = sub_coeff;
-  }
-
   for (int i = 0; i < sub_coeff_w * sub_coeff_h; i++) {
-    if (coeff_ptr[i] != 0) {
+    if (coeff[i] != 0) {
       found = 1;
       break;
     }
@@ -352,7 +334,7 @@ static INLINE double get_coeff_cabac_cost(
   if(!tr_skip) {
     uvg_encode_coeff_nxn((encoder_state_t*) state,
                          &cabac_copy,
-                         coeff_ptr,
+                         coeff,
                          cu_loc,
                          color,
                          scan_mode,
@@ -362,7 +344,7 @@ static INLINE double get_coeff_cabac_cost(
   else {
     uvg_encode_ts_residual((encoder_state_t* const)state,
       &cabac_copy,
-      coeff_ptr,
+      coeff,
       width,
       height,
       color,
@@ -426,6 +408,24 @@ double uvg_get_coeff_cost(
 
   const int width  = color == COLOR_Y ? cu_loc->width  : cu_loc->chroma_width;
   const int height = color == COLOR_Y ? cu_loc->height : cu_loc->chroma_height;
+  int x_local = cu_loc->x % LCU_WIDTH;
+  int y_local = cu_loc->y % LCU_WIDTH;
+  const int sub_coeff_w = color == COLOR_Y ? cu_loc->width : cu_loc->chroma_width;
+  const int sub_coeff_h = color == COLOR_Y ? cu_loc->height : cu_loc->chroma_height;
+  const int lcu_width = color == COLOR_Y ? LCU_WIDTH : LCU_WIDTH_C;
+
+
+  const coeff_t* coeff_ptr = NULL;
+  coeff_t sub_coeff[TR_MAX_WIDTH * TR_MAX_WIDTH];
+
+  if (coeff_order == COEFF_ORDER_LINEAR) {
+    coeff_ptr = coeff;
+  }
+  else {
+    // Coeff order CU
+    uvg_get_sub_coeff(sub_coeff, coeff, x_local, y_local, sub_coeff_w, sub_coeff_h, lcu_width);
+    coeff_ptr = sub_coeff;
+  }
 
   if (state->qp < state->encoder_control->cfg.fast_residual_cost_limit &&
       state->qp < MAX_FAST_COEFF_COST_QP && !tr_skip) {
@@ -437,15 +437,15 @@ double uvg_get_coeff_cost(
       return UINT32_MAX; // Hush little compiler don't you cry, not really gonna return anything after assert(0)
     } else {
       uint64_t weights = uvg_fast_coeff_get_weights(state);
-      uint32_t fast_cost = uvg_fast_coeff_cost(coeff, width, height, weights);
+      uint32_t fast_cost = uvg_fast_coeff_cost(coeff_ptr, width, height, weights);
       if (check_accuracy) {
-        double ccc = get_coeff_cabac_cost(state, coeff, cu_loc, color, scan_mode, tr_skip, cur_tu, coeff_order);
+        double ccc = get_coeff_cabac_cost(state, coeff_ptr, cu_loc, color, scan_mode, tr_skip, cur_tu);
         save_accuracy(state->qp, ccc, fast_cost);
       }
       return fast_cost;
     }
   } else {
-    double ccc = get_coeff_cabac_cost(state, coeff, cu_loc, color, scan_mode, tr_skip, cur_tu, coeff_order);
+    double ccc = get_coeff_cabac_cost(state, coeff_ptr, cu_loc, color, scan_mode, tr_skip, cur_tu);
     if (save_cccs) {
       save_ccc(state->qp, coeff, width * width, ccc);
     }