From d16d6e3dd8bb3ff4bd6f757952b39161886887fd Mon Sep 17 00:00:00 2001
From: Joose Sainio <joose.sainio@tuni.fi>
Date: Mon, 27 Jun 2022 10:36:44 +0300
Subject: [PATCH] [dual-tree] [lfnst] allow counting lfnst bits for chroma in
 dual-tree

---
 src/search_inter.c | 3 ++-
 src/search_intra.c | 3 ++-
 src/transform.c    | 5 +++--
 src/transform.h    | 3 ++-
 4 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/search_inter.c b/src/search_inter.c
index 472cbb43..c20f6db0 100644
--- a/src/search_inter.c
+++ b/src/search_inter.c
@@ -2186,7 +2186,8 @@ void uvg_cu_cost_inter_rd2(encoder_state_t * const state,
       v_pred,
       u_resi,
       v_resi,
-      &chorma_ts_out);
+      &chorma_ts_out,
+      UVG_BOTH_T);
     cbf_clear(&cur_cu->cbf, depth, COLOR_U);
     cbf_clear(&cur_cu->cbf, depth, COLOR_V);
     if (chorma_ts_out.best_u_cost + chorma_ts_out.best_v_cost < chorma_ts_out.best_combined_cost) {
diff --git a/src/search_intra.c b/src/search_intra.c
index 9b13aade..afb45602 100644
--- a/src/search_intra.c
+++ b/src/search_intra.c
@@ -1532,7 +1532,8 @@ int8_t uvg_search_intra_chroma_rdo(
             v_pred,
             u_resi,
             v_resi,
-            &chorma_ts_out);
+            &chorma_ts_out,
+            tree_type);
 
           // LFNST constraint failed
           if(chorma_ts_out.best_u_index == -1 && chorma_ts_out.best_combined_index == -1) {
diff --git a/src/transform.c b/src/transform.c
index 14fc9723..5d9748c7 100644
--- a/src/transform.c
+++ b/src/transform.c
@@ -482,7 +482,8 @@ void uvg_chroma_transform_search(
   uvg_pixel v_pred[1024],
   int16_t u_resi[1024],
   int16_t v_resi[1024],
-  uvg_chorma_ts_out_t* chorma_ts_out)
+  uvg_chorma_ts_out_t* chorma_ts_out,
+  enum uvg_tree_type tree_type)
 {
   ALIGNED(64) coeff_t u_coeff[LCU_WIDTH_C * LCU_WIDTH_C * 5];
   ALIGNED(64) uint8_t u_recon[LCU_WIDTH_C * LCU_WIDTH_C * 5];
@@ -699,7 +700,7 @@ void uvg_chroma_transform_search(
         scan_order,
         transforms[i] == CHROMA_TS);
     }
-    if(depth == 4 && state->encoder_control->cfg.lfnst && 0) {
+    if((depth == 4 || tree_type == UVG_CHROMA_T) && state->encoder_control->cfg.lfnst && 0) {
       if(uvg_is_lfnst_allowed(state, pred_cu, width, height, 0, 0 , UVG_CHROMA_T)) {
         const int lfnst_idx = pred_cu->cr_lfnst_idx;
         CABAC_FBITS_UPDATE(
diff --git a/src/transform.h b/src/transform.h
index 56d2ecd3..e8b496f3 100644
--- a/src/transform.h
+++ b/src/transform.h
@@ -110,7 +110,8 @@ void uvg_chroma_transform_search(
   uvg_pixel v_pred[1024],
   int16_t u_resi[1024],
   int16_t v_resi[1024],
-  uvg_chorma_ts_out_t* chorma_ts_out);
+  uvg_chorma_ts_out_t* chorma_ts_out,
+  enum uvg_tree_type tree_type);
 
 enum uvg_chroma_transforms {
   DCT7_CHROMA = 0,