[lfnst] cost on chroma when necessary and fixes

src/search.c

@@ -623,7 +623,7 @@ static double cu_rd_cost_tr_split_accurate(const encoder_state_t* const state,
   }
 
   if (uvg_is_lfnst_allowed(state, tr_cu, depth == 4 ? COLOR_UV : COLOR_Y, width, width, x_px, y_px)) {
-    const int lfnst_idx = tr_cu->lfnst_idx;
+    const int lfnst_idx = depth != 4 ? tr_cu->lfnst_idx : tr_cu->cr_lfnst_idx;
     CABAC_FBITS_UPDATE(
       cabac,
       &cabac->ctx.lfnst_idx_model[tr_cu->depth == 4],
@@ -941,6 +941,8 @@ static double search_cu(encoder_state_t * const state, int x, int y, int depth,
             lcu);
           intra_cost += uvg_cu_rd_cost_chroma(state, x_local, y_local, depth, &intra_search.pred_cu, lcu);
           intra_search.pred_cu.intra.mode = intra_mode;
+          intra_search.pred_cu.violates_lfnst_constrained_chroma = false;
+          intra_search.pred_cu.lfnst_last_scan_pos = false;
         }
 
       }

src/search_intra.c

@@ -1439,7 +1439,7 @@ int8_t uvg_search_intra_chroma_rdo(
         lfnst_modes_to_check[i] = i;
       }
     }
-    if(chroma_data->pred_cu.lfnst_idx) {
+    else if(chroma_data->pred_cu.lfnst_idx) {
       lfnst_modes_to_check[0] = chroma_data->pred_cu.lfnst_idx;
       lfnst_modes_to_check[1] = -1;
       lfnst_modes_to_check[2] = -1;
@@ -1524,10 +1524,8 @@ int8_t uvg_search_intra_chroma_rdo(
 
           if(chorma_ts_out.best_u_cost + chorma_ts_out.best_v_cost < chorma_ts_out.best_combined_cost) {
             chroma_data[mode_i].lfnst_costs[lfnst] += chorma_ts_out.best_u_cost + chorma_ts_out.best_v_cost;
-            if(lfnst == lfnst_modes_to_check[0] 
+            if( chroma_data[mode_i].lfnst_costs[lfnst] 
-              || (lfnst == 0 
+                < chroma_data[mode_i].lfnst_costs[best_lfnst_index] || lfnst_i == 0) {
-                && chroma_data[mode_i].lfnst_costs[lfnst] 
-                < chroma_data[mode_i].lfnst_costs[lfnst_modes_to_check[0]])) {
               chroma_data[mode_i].pred_cu.joint_cb_cr = 0;
               chroma_data[mode_i].pred_cu.tr_skip &= 1;
               chroma_data[mode_i].pred_cu.tr_skip |= (chorma_ts_out.best_u_index == CHROMA_TS) << COLOR_U;
@@ -1538,10 +1536,8 @@ int8_t uvg_search_intra_chroma_rdo(
           }
           else {
             chroma_data[mode_i].lfnst_costs[lfnst] += chorma_ts_out.best_combined_cost;
-            if (lfnst == lfnst_modes_to_check[0]
+            if (chroma_data[mode_i].lfnst_costs[lfnst]
-              || (lfnst == 0
+              < chroma_data[mode_i].lfnst_costs[best_lfnst_index] || lfnst_i == 0) {
-                && chroma_data[mode_i].lfnst_costs[lfnst]
-                < chroma_data[mode_i].lfnst_costs[lfnst_modes_to_check[0]])) {
               chroma_data[mode_i].pred_cu.joint_cb_cr = chorma_ts_out.best_combined_index;
               chroma_data[mode_i].pred_cu.tr_skip &= 1;
               best_lfnst_index = lfnst;

src/strategies/generic/quant-generic.c

@@ -499,7 +499,7 @@ int uvg_quantize_residual_generic(encoder_state_t *const state,
   } else {
   
     uvg_quant(state, coeff, coeff_out, width, width, color,
-      scan_order, cur_cu->type, cur_cu->tr_idx == MTS_SKIP && color == COLOR_Y, cur_cu->lfnst_idx);
+      scan_order, cur_cu->type, cur_cu->tr_idx == MTS_SKIP && color == COLOR_Y, lfnst_index);
   }
 
   // Check if there are any non-zero coefficients.

src/transform.c

@@ -32,6 +32,7 @@
 
 #include "transform.h"
 
+#include "encode_coding_tree.h"
 #include "image.h"
 #include "uvg266.h"
 #include "lfnst_tables.h"
@@ -537,6 +538,12 @@ void uvg_chroma_transform_search(
       uvg_get_scan_order(pred_cu->type, mode, depth);
     bool u_has_coeffs = false;
     bool v_has_coeffs = false;
+    if(pred_cu->cr_lfnst_idx) {
+      uvg_fwd_lfnst(pred_cu, width, height, COLOR_U, pred_cu->cr_lfnst_idx, &u_coeff[i * trans_offset]);
+      if (!IS_JCCR_MODE(transforms[i])) {
+        uvg_fwd_lfnst(pred_cu, width, height, COLOR_V, pred_cu->cr_lfnst_idx, &v_coeff[i * trans_offset]);
+      }
+    }
     quantize_chroma(
       state,
       depth,
@@ -556,19 +563,23 @@ void uvg_chroma_transform_search(
     
     if(pred_cu->type == CU_INTRA && transforms[i] != CHROMA_TS && depth == 4) {
       bool constraints[2] = { false, false };
-      uvg_derive_lfnst_constraints(pred_cu, depth, constraints, &u_coeff[i * trans_offset], width, height);
+      uvg_derive_lfnst_constraints(pred_cu, depth, constraints, u_quant_coeff, width, height);
       if(!IS_JCCR_MODE(transforms[i])) {
-        uvg_derive_lfnst_constraints(pred_cu, depth, constraints, &v_coeff[i * trans_offset], width, height);        
+        uvg_derive_lfnst_constraints(pred_cu, depth, constraints, v_quant_coeff, width, height);
       }
-      if ((constraints[0] || !constraints[1]) && pred_cu->cr_lfnst_idx != 0) continue;
+      if (!constraints[1] && (u_has_coeffs || v_has_coeffs) && pred_cu->cr_lfnst_idx != 0) continue;
     }
 
     if (IS_JCCR_MODE(transforms[i]) && !u_has_coeffs) continue;
 
     if (u_has_coeffs) {
+
       uvg_dequant(state, u_quant_coeff, &u_coeff[i * trans_offset], width, width, transforms[i] != JCCR_1 ? COLOR_U : COLOR_V,
         pred_cu->type, transforms[i] == CHROMA_TS);
       if (transforms[i] != CHROMA_TS) {
+        if (pred_cu->cr_lfnst_idx) {
+          uvg_inv_lfnst(pred_cu, width, height, COLOR_U, pred_cu->cr_lfnst_idx, &u_coeff[i * trans_offset]);
+        }
         uvg_itransform2d(state->encoder_control, u_recon_resi, &u_coeff[i * trans_offset], width,
           transforms[i] != JCCR_1 ? COLOR_U : COLOR_V, pred_cu);
       }
@@ -593,6 +604,9 @@ void uvg_chroma_transform_search(
       uvg_dequant(state, v_quant_coeff, &v_coeff[i * trans_offset], width, width, COLOR_V,
         pred_cu->type, transforms[i] == CHROMA_TS);
       if (transforms[i] != CHROMA_TS) {
+        if (pred_cu->cr_lfnst_idx) {
+          uvg_inv_lfnst(pred_cu, width, height, COLOR_V, pred_cu->cr_lfnst_idx, &v_coeff[i * trans_offset]);
+        }
         uvg_itransform2d(state->encoder_control, v_recon_resi, &v_coeff[i * trans_offset], width,
           transforms[i] != JCCR_1 ? COLOR_U : COLOR_V, pred_cu);
       }
@@ -663,7 +677,7 @@ void uvg_chroma_transform_search(
       double coeff_cost = uvg_get_coeff_cost(
         state,
         u_quant_coeff,
-        NULL,
+        pred_cu,
         width,
         COLOR_U,
         scan_order,
@@ -679,12 +693,33 @@ void uvg_chroma_transform_search(
       v_bits += uvg_get_coeff_cost(
         state,
         v_quant_coeff,
-        NULL,
+        pred_cu,
         width,
         COLOR_V,
         scan_order,
         transforms[i] == CHROMA_TS);
     }
+    if(depth == 4 && state->encoder_control->cfg.lfnst && 0) {
+      if(uvg_is_lfnst_allowed(state, pred_cu, COLOR_UV, width, height, 0 ,0)) {
+        const int lfnst_idx = pred_cu->cr_lfnst_idx;
+        CABAC_FBITS_UPDATE(
+          &state->search_cabac,
+          &state->search_cabac.ctx.lfnst_idx_model[1],
+          lfnst_idx != 0,
+          v_bits,
+          "lfnst_idx");
+        if (lfnst_idx > 0) {
+          CABAC_FBITS_UPDATE(
+            &state->search_cabac,
+            &state->search_cabac.ctx.lfnst_idx_model[2],
+            lfnst_idx == 2,
+            v_bits,
+            "lfnst_idx");
+        }
+      }
+      pred_cu->lfnst_last_scan_pos = false;
+      pred_cu->violates_lfnst_constrained_chroma = false;
+    }
     if (!IS_JCCR_MODE(transforms[i])) {
       double u_cost = UVG_CHROMA_MULT * ssd_u + u_bits * state->frame->lambda;
       double v_cost = UVG_CHROMA_MULT * ssd_v + v_bits * state->frame->lambda;