From cdc9b4960da8ba223578776474c058f061dbc567 Mon Sep 17 00:00:00 2001
From: Joose Sainio <joose.sainio@tuni.fi>
Date: Thu, 4 Jan 2024 13:46:19 +0200
Subject: [PATCH] fix

---
 src/search.c       | 2 +-
 src/search_intra.c | 2 +-
 src/transform.c    | 4 ++--
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/search.c b/src/search.c
index c353a914..cfc5f1da 100644
--- a/src/search.c
+++ b/src/search.c
@@ -1516,7 +1516,7 @@ static double search_cu(
           chroma_loc,
           tree_type);
       } else {
-        assert(cur_cu->cr_lfnst_idx == 0 && "If we don't have separate tree chroma lfnst index must be 0");
+        // assert(cur_cu->cr_lfnst_idx == 0 && "If we don't have separate tree chroma lfnst index must be 0");
       }
 
       // Set isp split cbfs here
diff --git a/src/search_intra.c b/src/search_intra.c
index a644ed9c..66a79901 100644
--- a/src/search_intra.c
+++ b/src/search_intra.c
@@ -1481,7 +1481,7 @@ int8_t uvg_search_intra_chroma_rdo(
       }
     }
     else {
-      lfnst_modes_to_check[0] = 0;
+      lfnst_modes_to_check[0] = chroma_data->pred_cu.lfnst_idx;
       lfnst_modes_to_check[1] = -1;
       lfnst_modes_to_check[2] = -1;
     }
diff --git a/src/transform.c b/src/transform.c
index 98728da0..2f4164ac 100644
--- a/src/transform.c
+++ b/src/transform.c
@@ -669,7 +669,7 @@ void uvg_chroma_transform_search(
       &u_coeff_cost,
       &v_coeff_cost);
     pred_cu->joint_cb_cr = old_jccr;
-    if (pred_cu->cr_lfnst_idx != 0 && !u_has_coeffs && !v_has_coeffs) goto reset_cabac;
+    if (pred_cu->cr_lfnst_idx != 0 && !u_has_coeffs && !v_has_coeffs && tree_type == UVG_CHROMA_T) goto reset_cabac;
     
     if(pred_cu->type == CU_INTRA && transforms[i] != CHROMA_TS && tree_type == UVG_CHROMA_T) {
       bool constraints[2] = { false, false };
@@ -677,7 +677,7 @@ void uvg_chroma_transform_search(
       if(!is_jccr) {
         uvg_derive_lfnst_constraints(pred_cu, constraints, v_quant_coeff, width, height, NULL, COLOR_V);
       }
-      if (!constraints[1] && (u_has_coeffs || v_has_coeffs) && pred_cu->cr_lfnst_idx != 0) goto reset_cabac;
+      if (!constraints[1] && (u_has_coeffs || v_has_coeffs) && pred_cu->cr_lfnst_idx != 0 && tree_type == UVG_CHROMA_T) goto reset_cabac;
     }
 
     if (is_jccr && !u_has_coeffs) goto reset_cabac;