From 88c33c048969fe7b067bc7e7150bba39bbe841fd Mon Sep 17 00:00:00 2001
From: Joose Sainio <joose.sainio@tuni.fi>
Date: Mon, 5 Sep 2022 16:51:23 +0300
Subject: [PATCH] [lfnst] Fix lfnst constraint checking for the new coeff order

---
 src/search_intra.c | 30 +++++++++++++++---------------
 src/transform.c    | 13 ++++++++++---
 src/transform.h    |  4 +++-
 3 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/src/search_intra.c b/src/search_intra.c
index 46db2f80..d5e8574b 100644
--- a/src/search_intra.c
+++ b/src/search_intra.c
@@ -205,13 +205,14 @@ static void derive_mts_constraints(cu_info_t *const pred_cu,
   const uint32_t * const scan = uvg_get_scan_order_table(SCAN_GROUP_4X4, scan_idx, log2_block_width, log2_block_height);
   const uint32_t * const scan_cg = uvg_get_scan_order_table(SCAN_GROUP_UNGROUPED, scan_idx, log2_block_width, log2_block_height);
 
-  const coeff_t* coeff = &lcu->coeff.y[xy_to_zorder(LCU_WIDTH, lcu_px.x, lcu_px.y)];
+  coeff_t coeff_y[TR_MAX_WIDTH * TR_MAX_WIDTH];
+  uvg_get_sub_coeff(coeff_y, lcu->coeff.y, lcu_px.x, lcu_px.y, width, height, LCU_WIDTH);
 
   signed scan_cg_last = -1;
   signed scan_pos_last = -1;
 
   for (int i = 0; i < width * height; i++) {
-    if (coeff[scan[i]]) {
+    if (coeff_y[scan[i]]) {
       scan_pos_last = i;
       sig_coeffgroup_flag[scan_cg[i >> log2_cg_size]] = 1;
     }
@@ -420,10 +421,11 @@ static double search_intra_trdepth(
             pred_cu,
             depth,
             constraints,
-            &lcu->coeff.y[scan_offset],
+            lcu->coeff.y,
             width,
-            height
-            );
+            height,
+            &lcu_px,
+            COLOR_Y);
         }
 
         if (!constraints[1] && cbf_is_set(pred_cu->cbf, depth, COLOR_Y)) {
@@ -493,10 +495,6 @@ static double search_intra_trdepth(
         pred_cu->intra.mode_chroma = chroma_mode;
         pred_cu->joint_cb_cr = 4;
         // TODO: Maybe check the jccr mode here also but holy shit is the interface of search_intra_rdo bad currently
-        const unsigned scan_offset = xy_to_zorder(
-          LCU_WIDTH_C,
-          lcu_px.x,
-          lcu_px.y);
         uvg_intra_recon_cu(
           state,
           x_px,
@@ -526,10 +524,11 @@ static double search_intra_trdepth(
             pred_cu,
             depth,
             constraints,
-            &lcu->coeff.u[scan_offset],
+            lcu->coeff.u,
             width_c,
-            width_c
-            );
+            width_c,
+            &lcu_px,
+            COLOR_U);
           if (constraints[0] || !constraints[1]) {
             best_lfnst_idx = 0;
             continue;
@@ -538,10 +537,11 @@ static double search_intra_trdepth(
             pred_cu,
             depth,
             constraints,
-            &lcu->coeff.u[scan_offset],
+            lcu->coeff.u,
             width_c,
-            width_c
-            );
+            width_c,
+            &lcu_px,
+            COLOR_U);
           if (constraints[0] || !constraints[1]) {
             best_lfnst_idx = 0;
             continue;
diff --git a/src/transform.c b/src/transform.c
index 825df5b2..8b903579 100644
--- a/src/transform.c
+++ b/src/transform.c
@@ -178,7 +178,9 @@ void uvg_derive_lfnst_constraints(
   bool* constraints,
   const coeff_t* coeff,
   const int width,
-  const int height)
+  const int height,
+  const vector2d_t * const lcu_px,
+  color_t color)
 {
   coeff_scan_order_t scan_idx = uvg_get_scan_order(pred_cu->type, pred_cu->intra.mode, depth);
   // ToDo: large block support in VVC?
@@ -187,6 +189,11 @@ void uvg_derive_lfnst_constraints(
   const uint32_t* scan = uvg_g_sig_last_scan[scan_idx][log2_block_size - 1];
 
   signed scan_pos_last = -1;
+  coeff_t temp[TR_MAX_WIDTH * TR_MAX_WIDTH];
+  if(lcu_px != NULL) {
+    uvg_get_sub_coeff(temp, coeff, lcu_px->x, lcu_px->y, width, height, color == COLOR_Y? LCU_WIDTH : LCU_WIDTH_C);
+    coeff = temp;
+  }
 
   for (int i = 0; i < width * height; i++) {
     if (coeff[scan[i]]) {
@@ -576,9 +583,9 @@ void uvg_chroma_transform_search(
     
     if(pred_cu->type == CU_INTRA && transforms[i] != CHROMA_TS && (depth == 4 || tree_type == UVG_CHROMA_T)) {
       bool constraints[2] = { false, false };
-      uvg_derive_lfnst_constraints(pred_cu, depth, constraints, u_quant_coeff, width, height);
+      uvg_derive_lfnst_constraints(pred_cu, depth, constraints, u_quant_coeff, width, height, NULL, COLOR_U);
       if(!IS_JCCR_MODE(transforms[i])) {
-        uvg_derive_lfnst_constraints(pred_cu, depth, constraints, v_quant_coeff, width, height);
+        uvg_derive_lfnst_constraints(pred_cu, depth, constraints, v_quant_coeff, width, height, NULL, COLOR_V);
       }
       if (!constraints[1] && (u_has_coeffs || v_has_coeffs) && pred_cu->cr_lfnst_idx != 0) continue;
     }
diff --git a/src/transform.h b/src/transform.h
index 78a2325a..6fdef411 100644
--- a/src/transform.h
+++ b/src/transform.h
@@ -78,7 +78,9 @@ void uvg_derive_lfnst_constraints(
   bool* constraints,
   const coeff_t* coeff,
   const int width,
-  const int height);
+  const int height,
+  const vector2d_t * const ,
+  color_t color);
 
 typedef struct {
   double best_u_cost;