From ed8496e57e6010082ce1920b7286aa35830ed7f7 Mon Sep 17 00:00:00 2001
From: Joose Sainio <joose.sainio@tuni.fi>
Date: Wed, 8 Jun 2022 15:06:40 +0300
Subject: [PATCH] [dual-tree] Matches except for cutoff bottom CTUs for forced
 depth=1

---
 src/intra.c  | 24 ++++++++++++++++++++++++
 src/intra.h  |  9 +++++++++
 src/search.c |  3 +++
 3 files changed, 36 insertions(+)

diff --git a/src/intra.c b/src/intra.c
index 01e3eccd..df9662b5 100644
--- a/src/intra.c
+++ b/src/intra.c
@@ -1417,6 +1417,30 @@ void uvg_intra_predict(
   }
 }
 
+// This function works on luma coordinates 
+const cu_info_t* kvz_get_co_located_luma_cu(
+  int x,
+  int y,
+  int width,
+  int height,
+  const lcu_t* const lcu,
+  const cu_array_t* const cu_array,
+  enum kvz_tree_type tree_type)
+{
+  assert((cu_array || lcu) && !(cu_array && lcu));
+  assert(tree_type != KVZ_LUMA_T && "Luma only CU shouldn't need colocated luma CU");
+  if(tree_type == KVZ_CHROMA_T) {
+    x += width >> 1;
+    y += height >> 1;
+  }
+  if(cu_array) {
+    return uvg_cu_array_at_const(cu_array, x, y);
+  }
+  else {
+    return LCU_GET_CU_AT_PX(lcu, SUB_SCU(x), SUB_SCU(y));
+  }
+}
+
 
 static void intra_recon_tb_leaf(
   encoder_state_t* const state,
diff --git a/src/intra.h b/src/intra.h
index 6357bc29..3d48e16c 100644
--- a/src/intra.h
+++ b/src/intra.h
@@ -146,4 +146,13 @@ void uvg_intra_recon_cu(
   lcu_t *lcu, 
   enum kvz_tree_type tree_type);
 
+const cu_info_t* kvz_get_co_located_luma_cu(
+  int x,
+  int y,
+  int width,
+  int height,
+  const lcu_t* const lcu,
+  const cu_array_t* const cu_array,
+  enum kvz_tree_type tree_type);
+
 int uvg_get_mip_flag_context(int x, int y, int width, int height, const lcu_t* lcu, cu_array_t* const cu_a);
diff --git a/src/search.c b/src/search.c
index cec50a97..f94ca0fd 100644
--- a/src/search.c
+++ b/src/search.c
@@ -944,6 +944,9 @@ static double search_cu(
           // 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.
+          if(tree_type == KVZ_CHROMA_T) {
+            intra_search.pred_cu.intra = kvz_get_co_located_luma_cu(x, y, luma_width, luma_width, lcu, NULL, KVZ_CHROMA_T)->intra;
+          }
           intra_search.pred_cu.intra.mode_chroma = intra_search.pred_cu.intra.mode;
           if (ctrl->cfg.rdo >= 3 || ctrl->cfg.jccr || ctrl->cfg.lfnst) {
             uvg_search_cu_intra_chroma(state, x, y, depth, lcu, &intra_search, tree_type);