From 85f6b0039449ef62810b411e56cc1cd83be09dd8 Mon Sep 17 00:00:00 2001
From: siivonek <kari.siivonen@tuni.fi>
Date: Mon, 26 Sep 2022 14:44:51 +0300
Subject: [PATCH] [isp] Add lfnst asserts. Fix error in MTS search. Fix chroma
 lfnst index when no coefficients present.

---
 src/encode_coding_tree.c | 6 ++++++
 src/search_intra.c       | 8 +++-----
 src/transform.c          | 1 +
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/src/encode_coding_tree.c b/src/encode_coding_tree.c
index c2be1395..a739cebf 100644
--- a/src/encode_coding_tree.c
+++ b/src/encode_coding_tree.c
@@ -230,6 +230,12 @@ static bool encode_lfnst_idx(
     return true;
   }
   else {
+    if(depth != 4 || color == COLOR_Y) {
+      assert(pred_cu->lfnst_idx == 0);
+    }
+    if(depth == 4 && color != COLOR_Y) {
+      assert(pred_cu->cr_lfnst_idx == 0);
+    }
     return false;
   }
 }
diff --git a/src/search_intra.c b/src/search_intra.c
index 224482dc..03609552 100644
--- a/src/search_intra.c
+++ b/src/search_intra.c
@@ -343,9 +343,7 @@ static double search_intra_trdepth(
     
     const int max_tb_size = TR_MAX_WIDTH;
     // LFNST search params
-    int max_lfnst_idx = width > max_tb_size || height > max_tb_size ?
-                                0 :
-                                2;
+    int max_lfnst_idx = width > max_tb_size || height > max_tb_size ? 0 : 2;
     if(pred_cu->intra.mip_flag && (width < 16 || height < 16)) {
       max_lfnst_idx = 0;
     }
@@ -379,8 +377,8 @@ static double search_intra_trdepth(
           pred_cu->mts_last_scan_pos = 0;
           pred_cu->violates_mts_coeff_constraint = 0;
 
-          if ((trafo == MTS_SKIP && width > (1 << state->encoder_control->cfg.trskip_max_size))
-            || !state->encoder_control->cfg.trskip_enable) {
+          if (trafo == MTS_SKIP && (width > (1 << state->encoder_control->cfg.trskip_max_size)
+            || !state->encoder_control->cfg.trskip_enable)) {
             continue;
           }
         }
diff --git a/src/transform.c b/src/transform.c
index ee9e79ec..954d7836 100644
--- a/src/transform.c
+++ b/src/transform.c
@@ -580,6 +580,7 @@ void uvg_chroma_transform_search(
       &u_has_coeffs,
       &v_has_coeffs,
       pred_cu->cr_lfnst_idx);
+      if(pred_cu->cr_lfnst_idx !=0 && !u_has_coeffs && !v_has_coeffs) continue;
     
     if(pred_cu->type == CU_INTRA && transforms[i] != CHROMA_TS && (depth == 4 || tree_type == UVG_CHROMA_T)) {
       bool constraints[2] = { false, false };