From 3a73abd26465b14226a16caf49da7419ef1e4544 Mon Sep 17 00:00:00 2001
From: Joose Sainio <joose.sainio@tuni.fi>
Date: Thu, 19 Aug 2021 14:27:55 +0300
Subject: [PATCH] =?UTF-8?q?[jccr]=20disable=20jccr=20for=20blocks=20when?=
 =?UTF-8?q?=20tr-depth=20!=3D=20depth,=20i.e.=2064=C3=9764?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/encode_coding_tree.c |  2 +-
 src/intra.c              |  2 +-
 src/search.c             | 25 +++++++++++++++++++++----
 src/transform.c          |  2 +-
 4 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/src/encode_coding_tree.c b/src/encode_coding_tree.c
index 316330d6..a25304db 100644
--- a/src/encode_coding_tree.c
+++ b/src/encode_coding_tree.c
@@ -429,7 +429,7 @@ static void encode_transform_unit(encoder_state_t * const state,
 
   bool chroma_cbf_set = cbf_is_set(cur_pu->cbf, depth, COLOR_U) ||
                         cbf_is_set(cur_pu->cbf, depth, COLOR_V);
-  if (chroma_cbf_set) {
+  if (chroma_cbf_set || joint_chroma) {
     encode_chroma_tu(state, x, y, depth, width_c, cur_pu, &scan_idx, coeff, joint_chroma);
   }
 }
diff --git a/src/intra.c b/src/intra.c
index 64b1e643..d5924c01 100644
--- a/src/intra.c
+++ b/src/intra.c
@@ -619,7 +619,7 @@ static void intra_recon_tb_leaf(
   }
 
   kvz_pixels_blit(pred, block , width, width, width, lcu_width);
-  if(color != COLOR_Y) {
+  if(color != COLOR_Y && cfg->jccr) {
     kvz_pixels_blit(pred, block2, width, width, width, lcu_width);
   }
 }
diff --git a/src/search.c b/src/search.c
index 5c5f481d..67788855 100644
--- a/src/search.c
+++ b/src/search.c
@@ -310,6 +310,7 @@ double kvz_cu_rd_cost_chroma(const encoder_state_t *const state,
   cu_info_t *const tr_cu = LCU_GET_CU_AT_PX(lcu, x_px, y_px);
 
   double tr_tree_bits = 0;
+  double joint_cbcr_tr_tree_bits = 0;
   double coeff_bits = 0;
   double joint_coeff_bits = 0;
 
@@ -328,13 +329,21 @@ double kvz_cu_rd_cost_chroma(const encoder_state_t *const state,
     if (tr_depth == 0 || cbf_is_set(pred_cu->cbf, depth - 1, COLOR_U)) {
       tr_tree_bits += CTX_ENTROPY_FBITS(ctx, cbf_is_set(pred_cu->cbf, depth, COLOR_U));
     }
+    if(state->encoder_control->cfg.jccr) {
+      joint_cbcr_tr_tree_bits += CTX_ENTROPY_FBITS(ctx, pred_cu->joint_cb_cr & 1);
+    }
     int is_set = cbf_is_set(pred_cu->cbf, depth, COLOR_U);
     ctx = &(state->cabac.ctx.qt_cbf_model_cr[is_set]);
     if (tr_depth == 0 || cbf_is_set(pred_cu->cbf, depth - 1, COLOR_V)) {
       tr_tree_bits += CTX_ENTROPY_FBITS(ctx, cbf_is_set(pred_cu->cbf, depth, COLOR_V));
     }
+    if(state->encoder_control->cfg.jccr) {
+      ctx = &(state->cabac.ctx.qt_cbf_model_cr[pred_cu->joint_cb_cr & 1]);
+      joint_cbcr_tr_tree_bits += CTX_ENTROPY_FBITS(ctx, (pred_cu->joint_cb_cr & 2) >> 1);
+    }
   }
 
+
   if (tr_cu->tr_depth > depth) {
     int offset = LCU_WIDTH >> (depth + 1);
     int sum = 0;
@@ -347,6 +356,13 @@ double kvz_cu_rd_cost_chroma(const encoder_state_t *const state,
     return sum + tr_tree_bits * state->lambda;
   }
 
+  if (state->encoder_control->cfg.jccr) {
+    const cabac_ctx_t* ctx = &(state->cabac.ctx.joint_cb_cr[cbf_is_set(pred_cu->cbf, depth, COLOR_U) * 2 + cbf_is_set(pred_cu->cbf, depth, COLOR_V) - 1]);
+    tr_tree_bits += CTX_ENTROPY_FBITS(ctx, 0);
+    ctx = &(state->cabac.ctx.joint_cb_cr[(pred_cu->joint_cb_cr & 1) * 2 + ((pred_cu->joint_cb_cr & 2) >> 1) - 1]);
+    joint_cbcr_tr_tree_bits += CTX_ENTROPY_FBITS(ctx, 1);
+  }
+
   // Chroma SSD
   int ssd = 0;
   int joint_ssd = 0;
@@ -383,21 +399,22 @@ double kvz_cu_rd_cost_chroma(const encoder_state_t *const state,
     }
   }
 
+
   double bits = tr_tree_bits + coeff_bits;
-  double joint_bits = tr_tree_bits + joint_coeff_bits;
+  double joint_bits = joint_cbcr_tr_tree_bits + joint_coeff_bits;
 
   double cost = (double)ssd + bits * state->c_lambda;
   double joint_cost = (double)joint_ssd + joint_bits * state->c_lambda;
   if ((cost < joint_cost || !pred_cu->joint_cb_cr) || !state->encoder_control->cfg.jccr) {
-    tr_cu->joint_cb_cr = 0;
+    pred_cu->joint_cb_cr = 0;
     return cost;    
   }
   cbf_clear(&pred_cu->cbf, depth, COLOR_U);
   cbf_clear(&pred_cu->cbf, depth, COLOR_V);
-  if (tr_cu->joint_cb_cr & 1) {
+  if (pred_cu->joint_cb_cr & 1) {
     cbf_set(&pred_cu->cbf, depth, COLOR_U);
   }
-  if (tr_cu->joint_cb_cr & 2) {
+  if (pred_cu->joint_cb_cr & 2) {
     cbf_set(&pred_cu->cbf, depth, COLOR_V);
   }
   int lcu_width = LCU_WIDTH_C;
diff --git a/src/transform.c b/src/transform.c
index 094845b4..bb768218 100644
--- a/src/transform.c
+++ b/src/transform.c
@@ -539,7 +539,7 @@ void kvz_quantize_lcu_residual(encoder_state_t * const state,
     if (chroma) {
       quantize_tr_residual(state, COLOR_U, x, y, depth, cur_pu, lcu, early_skip);
       quantize_tr_residual(state, COLOR_V, x, y, depth, cur_pu, lcu, early_skip);
-      if(state->encoder_control->cfg.jccr){
+      if(state->encoder_control->cfg.jccr && cur_pu->tr_depth == cur_pu->depth){
         quantize_tr_residual(state, COLOR_UV, x, y, depth, cur_pu, lcu, early_skip);
       }
     }