[lfnst] various small fixes

src/encode_coding_tree.c

@@ -159,7 +159,7 @@ static bool can_use_lfnst_with_isp(const int width, const int height, const int
 }
 
  bool uvg_is_lfnst_allowed(
-  encoder_state_t* const state,
+  const encoder_state_t* const state,
   const cu_info_t* const pred_cu,
   const int color,
   const int width,

src/encode_coding_tree.h

@@ -42,7 +42,7 @@
 
 bool uvg_is_mts_allowed(const encoder_state_t* const state, cu_info_t* const pred_cu);
 bool uvg_is_lfnst_allowed(
-  encoder_state_t* const state,
+  const encoder_state_t* const state,
   const cu_info_t* const pred_cu,
   const int color,
   const int width,

src/search.c

@@ -898,7 +898,7 @@ static double search_cu(encoder_state_t * const state, int x, int y, int depth,
           // into account, so there is less of a chanse of luma mode being
           // really bad for chroma.
           intra_search.pred_cu.intra.mode_chroma = intra_search.pred_cu.intra.mode;
-          if (ctrl->cfg.rdo >= 3) {
+          if (ctrl->cfg.rdo >= 3 || ctrl->cfg.jccr || ctrl->cfg.lfnst) {
             uvg_search_cu_intra_chroma(state, x, y, depth, lcu, &intra_search);
 
             if (intra_search.pred_cu.joint_cb_cr == 0) {
@@ -912,12 +912,18 @@ static double search_cu(encoder_state_t * const state, int x, int y, int depth,
           else {
             intra_search.pred_cu.intra.mode_chroma = 0;
           }
+          uvg_intra_recon_cu(state,
+            x, y,
+            depth, &intra_search,
+            &intra_search.pred_cu,
+            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.intra.mode = intra_mode;
-        intra_cost += intra_search.cost;
+
       }
       if (intra_cost < cost) {
-        cost = intra_search.cost;
+        cost = intra_cost;
         *cur_cu = intra_search.pred_cu;
         cur_cu->type = CU_INTRA;
       }

src/search_intra.c

@@ -344,7 +344,7 @@ static double search_intra_trdepth(
                   tr_depth ?
                     max_lfnst_idx :
                     0;
-    for (int i = start_idx; i < end_idx; ++i) {
+    for (int i = start_idx; i < end_idx + 1; ++i) {
       search_data->lfnst_costs[i] = MAX_DOUBLE;
     }
     bool constraints[2] = { pred_cu->violates_lfnst_constrained_luma,
@@ -1588,8 +1588,8 @@ int8_t uvg_search_cu_intra_chroma(encoder_state_t * const state,
   // is always one of the modes, so 2 means the final decision is made
   // between luma mode and one other mode that looks the best
   // according to search_intra_chroma_rough.
-  const int8_t modes_in_depth[5] = { 1, 1, 1, 1, 2 };
-  int num_modes = modes_in_depth[depth];
+  // const int8_t modes_in_depth[5] = { 1, 1, 1, 1, 1 };
+  int num_modes = 5;
 
   if (state->encoder_control->cfg.rdo >= 3) {
     num_modes = total_modes;
@@ -1597,11 +1597,12 @@ int8_t uvg_search_cu_intra_chroma(encoder_state_t * const state,
 
   intra_search_data_t chroma_data[8];
   FILL(chroma_data, 0);
-  for (int i = 0; i < num_modes; i++) {
+  for (int i = state->encoder_control->cfg.rdo >= 3 ? 0 : 4; i < num_modes; i++) {
     chroma_data[i].pred_cu = *cur_pu;
     chroma_data[i].pred_cu.intra.mode_chroma = modes[i];
     chroma_data[i].pred_cu.intra.mode = -1;
     chroma_data[i].cost = 0;
+    memcpy(chroma_data[i].lfnst_costs, search_data->lfnst_costs, sizeof(double) * 3);
   }
   // Don't do rough mode search if all modes are selected.
   // FIXME: It might make more sense to only disable rough search if

src/transform.c

@@ -557,7 +557,7 @@ void uvg_chroma_transform_search(
       if(!IS_JCCR_MODE(transforms[i])) {
         uvg_derive_lfnst_constraints(pred_cu, depth, constraints, &v_coeff[i * trans_offset], width, height);        
       }
-      if (constraints[0] || !constraints[1] && pred_cu->lfnst_idx != 0) continue;
+      if ((constraints[0] || !constraints[1]) && pred_cu->lfnst_idx != 0) continue;
     }
 
     if (IS_JCCR_MODE(transforms[i]) && !u_has_coeffs) continue;