[intra] Perform intra chroma search before the mode is selected

src/encode_coding_tree.c

@@ -747,6 +747,7 @@ static void encode_transform_coeff(encoder_state_t * const state,
       && (depth != 4 || only_chroma) 
       && state->encoder_control->cfg.jccr
       ) {
+      assert(cur_pu->joint_cb_cr < 4 && "JointCbCr is in search state.");
       cabac->cur_ctx = &cabac->ctx.joint_cb_cr[cb_flag_u * 2 + cb_flag_v - 1];
       CABAC_BIN(cabac, cur_pu->joint_cb_cr != 0, "tu_joint_cbcr_residual_flag");
     }

src/intra.c

@@ -1576,7 +1576,7 @@ void uvg_intra_recon_cu(
     }
   } else {
     const bool has_luma = mode_luma != -1;
-    const bool has_chroma = mode_chroma != -1 &&  (x % 8 == 0 && y % 8 == 0);
+    const bool has_chroma = mode_chroma != -1 && (x % 8 == 0 && y % 8 == 0);
    
     // Process a leaf TU.
     if (has_luma) {
@@ -1588,7 +1588,7 @@ void uvg_intra_recon_cu(
     }
 
     uvg_quantize_lcu_residual(state, has_luma, has_chroma && !(search_data->pred_cu.joint_cb_cr & 3),
-      search_data->pred_cu.joint_cb_cr != 4 && state->encoder_control->cfg.jccr && (x % 8 == 0 && y % 8 == 0),
+      search_data->pred_cu.joint_cb_cr & 3 && state->encoder_control->cfg.jccr && has_chroma,
       x, y, depth, cur_cu, lcu, false);
   }
 }

src/search.c

@@ -854,7 +854,49 @@ static double search_cu(encoder_state_t * const state, int x, int y, int depth,
         intra_cost += pred_mode_type_bits * state->lambda;
       }
 #endif
-      if (intra_search.cost < cost) {
+      double intra_cost = intra_search.cost;
+      if (intra_cost < cost) {
+        int8_t intra_mode = intra_search.pred_cu.intra.mode;
+        if(state->encoder_control->cfg.cclm) {
+          intra_search.pred_cu.intra.mode_chroma = -1;
+          uvg_intra_recon_cu(state,
+            x, y,
+            depth, &intra_search,
+            &intra_search.pred_cu,
+            lcu);
+
+          downsample_cclm_rec(
+            state, x, y, cu_width / 2, cu_width / 2, lcu->rec.y, lcu->left_ref.y[64]
+          );
+        }
+        intra_search.pred_cu.joint_cb_cr = 0;
+
+        // TODO: This heavily relies to square CUs
+        if ((depth != 4 || (x % 8 && y % 8)) && state->encoder_control->chroma_format != KVZ_CSP_400) {
+          // There is almost no benefit to doing the chroma mode search for
+          // rd2. Possibly because the luma mode search already takes chroma
+          // 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) {
+            uvg_search_cu_intra_chroma(state, x, y, depth, lcu, &intra_search);
+
+            if (intra_search.pred_cu.joint_cb_cr == 0) {
+              intra_search.pred_cu.joint_cb_cr = 4;
+            }
+
+          }
+          else if (!intra_search.pred_cu.intra.mip_flag) {
+            intra_search.pred_cu.intra.mode_chroma = intra_search.pred_cu.intra.mode;
+          }
+          else {
+            intra_search.pred_cu.intra.mode_chroma = 0;
+          }
+        }
+        intra_search.pred_cu.intra.mode = intra_mode;
+        intra_cost += intra_search.cost;
+      }
+      if (intra_cost < cost) {
         cost = intra_search.cost;
         *cur_cu = intra_search.pred_cu;
         cur_cu->type = CU_INTRA;
@@ -866,53 +908,19 @@ static double search_cu(encoder_state_t * const state, int x, int y, int depth,
     if (cur_cu->type == CU_INTRA) {
       assert(cur_cu->part_size == SIZE_2Nx2N || cur_cu->part_size == SIZE_NxN);
 
-      intra_search.pred_cu.intra.mode_chroma = -1; // don't reconstruct chroma before search is performed for it
+      if ((depth == 4 && (x % 8 == 0 || y % 8 == 0)) || state->encoder_control->chroma_format == KVZ_CSP_400) {
+        intra_search.pred_cu.intra.mode_chroma = -1; 
+      }
       lcu_fill_cu_info(lcu, x_local, y_local, cu_width, cu_width, cur_cu);
       uvg_intra_recon_cu(state,
                          x, y,
                          depth, &intra_search,
                          NULL, 
                          lcu);
+      if (cur_cu->joint_cb_cr == 4) cur_cu->joint_cb_cr = 0;
+      lcu_fill_cu_info(lcu, x_local, y_local, cu_width, cu_width, cur_cu);
 
-      downsample_cclm_rec(
-        state, x, y, cu_width / 2, cu_width / 2, lcu->rec.y, lcu->left_ref.y[64]
-      );
-      cur_cu->joint_cb_cr = 0;
 
-      // TODO: This heavily relies to square CUs
-      if ((depth != 4 || (x % 8 && y % 8)) && state->encoder_control->chroma_format != UVG_CSP_400) {
-        // There is almost no benefit to doing the chroma mode search for
-        // rd2. Possibly because the luma mode search already takes chroma
-        // into account, so there is less of a chanse of luma mode being
-        // really bad for chroma.
-        intra_search.pred_cu.intra.mode_chroma = cur_cu->intra.mode; // skip luma
-        if (ctrl->cfg.rdo >= 3) {
-          cur_cu->intra.mode_chroma = uvg_search_cu_intra_chroma(state, x, y, depth, lcu, &intra_search);
-
-          if (intra_search.pred_cu.joint_cb_cr == 0) {
-            intra_search.pred_cu.joint_cb_cr = 4;
-            cur_cu->tr_skip |= intra_search.pred_cu.tr_skip;
-          }
-          else cur_cu->joint_cb_cr = intra_search.pred_cu.joint_cb_cr;
-
-          lcu_fill_cu_info(lcu, x_local, y_local, cu_width, cu_width, cur_cu);
-        }
-        else if(!cur_cu->intra.mip_flag) {
-          cur_cu->intra.mode_chroma = cur_cu->intra.mode;
-          intra_search.pred_cu.intra.mode_chroma = cur_cu->intra.mode;
-        }
-        else {
-          cur_cu->intra.mode_chroma = 0;
-          intra_search.pred_cu.intra.mode_chroma = 0;
-        }
-        intra_search.pred_cu.intra.mode = -1; // skip luma
-        uvg_intra_recon_cu(state,
-                           x, y, // TODO: as does this
-                           depth, &intra_search,
-                           NULL,
-                           lcu);
-
-      }
     } else if (cur_cu->type == CU_INTER) {
 
       if (!cur_cu->skipped) {

src/search_intra.c

@@ -1488,7 +1488,7 @@ int8_t uvg_search_cu_intra_chroma(encoder_state_t * const state,
 {
   const vector2d_t lcu_px = { SUB_SCU(x_px), SUB_SCU(y_px) };
 
-  cu_info_t *cur_pu = LCU_GET_CU_AT_PX(lcu, lcu_px.x, lcu_px.y);
+  const cu_info_t *cur_pu = &search_data->pred_cu;
   int8_t intra_mode = !cur_pu->intra.mip_flag ? cur_pu->intra.mode : 0;
   
   int8_t modes[8] = { 0, 50, 18, 1, intra_mode, 81, 82, 83 };