From 58c6af8c87ceeedd62e8e7b4d40aaa775f7c59e0 Mon Sep 17 00:00:00 2001
From: Joose Sainio <joose.sainio@tuni.fi>
Date: Mon, 19 Sep 2022 08:44:49 +0300
Subject: [PATCH] [mtt] Add function for easily getting all split cu_locs

---
 src/cu.c     | 39 +++++++++++++++++++++++++++++++++++++++
 src/cu.h     |  1 +
 src/search.c | 15 ++++++---------
 3 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/src/cu.c b/src/cu.c
index aedf341c..4ae74da0 100644
--- a/src/cu.c
+++ b/src/cu.c
@@ -315,3 +315,42 @@ void uvg_cu_loc_ctor(cu_loc_t* loc, int x, int y, int width, int height)
   loc->chroma_width = MAX(width >> 1, 4);
   loc->chroma_height = MAX(height >> 1, 4);
 }
+
+
+void uvg_get_split_locs(const cu_loc_t* const origin, enum split_type split, cu_loc_t out[4])
+{
+  const int half_width = origin->width >> 1;
+  const int half_height = origin->height >> 1;
+  const int quarter_width = origin->width >> 2;
+  const int quarter_height = origin->height >> 2;
+
+  switch (split) {
+    case NO_SPLIT:
+      assert(0 && "trying to get split from no split");
+    break;
+    case QT_SPLIT:
+      uvg_cu_loc_ctor(&out[0], origin->x, origin->y, half_width, half_height);
+      uvg_cu_loc_ctor(&out[1], origin->x + half_width, origin->y, half_width, half_height);
+      uvg_cu_loc_ctor(&out[2], origin->x, origin->y + half_height, half_width, half_height);
+      uvg_cu_loc_ctor(&out[3], origin->x + half_width, origin->y + half_height, half_width, half_height);
+    break;
+    case BT_HOR_SPLIT:
+      uvg_cu_loc_ctor(&out[0], origin->x, origin->y, origin->width, half_height);
+      uvg_cu_loc_ctor(&out[1], origin->x, origin->y + half_height, origin->width, half_height);
+    break;
+    case BT_VER_SPLIT:
+      uvg_cu_loc_ctor(&out[0], origin->x, origin->y, half_width, origin->height);
+      uvg_cu_loc_ctor(&out[1], origin->x + half_width, origin->y, half_width, origin->height);
+    break;
+    case TT_HOR_SPLIT:
+      uvg_cu_loc_ctor(&out[0], origin->x, origin->y, origin->width, quarter_height);
+      uvg_cu_loc_ctor(&out[1], origin->x, origin->y + quarter_height, origin->width, half_height);
+      uvg_cu_loc_ctor(&out[2], origin->x, origin->y + quarter_height + half_height, origin->width, quarter_height);
+    break;
+    case TT_VER_SPLIT:
+      uvg_cu_loc_ctor(&out[0], origin->x, origin->y, quarter_width, origin->height);
+      uvg_cu_loc_ctor(&out[1], origin->x + quarter_width, origin->y, half_width, origin->height);
+      uvg_cu_loc_ctor(&out[2], origin->x + quarter_width + half_width, origin->y, quarter_width, origin->height);
+    break;
+  }
+}
\ No newline at end of file
diff --git a/src/cu.h b/src/cu.h
index e05df7de..f9021b3c 100644
--- a/src/cu.h
+++ b/src/cu.h
@@ -185,6 +185,7 @@ typedef struct {
 
 void uvg_cu_loc_ctor(cu_loc_t *loc, int x, int y, int width, int height);
 
+void uvg_get_split_locs(const cu_loc_t* const origin, enum split_type split, cu_loc_t out[4]);
 
 #define CU_GET_MV_CAND(cu_info_ptr, reflist) \
   (((reflist) == 0) ? (cu_info_ptr)->inter.mv_cand0 : (cu_info_ptr)->inter.mv_cand1)
diff --git a/src/search.c b/src/search.c
index 8778d081..bcfd82bc 100644
--- a/src/search.c
+++ b/src/search.c
@@ -1336,22 +1336,19 @@ static double search_cu(
     // It is ok to interrupt the search as soon as it is known that
     // the split costs at least as much as not splitting.
     if (cur_cu->type == CU_NOTSET || cbf || state->encoder_control->cfg.cu_split_termination == UVG_CU_SPLIT_TERMINATION_OFF) {
-      cu_loc_t new_cu_loc;
+      cu_loc_t new_cu_loc[4];
+      uvg_get_split_locs(cu_loc, QT_SPLIT, new_cu_loc);
       if (split_cost < cost) {
-        uvg_cu_loc_ctor(&new_cu_loc, x, y, half_cu, half_cu);
-        split_cost += search_cu(state, &new_cu_loc, work_tree, tree_type, new_split);
+        split_cost += search_cu(state, &new_cu_loc[0], work_tree, tree_type, new_split);
       }
       if (split_cost < cost) {
-        uvg_cu_loc_ctor(&new_cu_loc, x + half_cu, y, half_cu, half_cu);
-        split_cost += search_cu(state, &new_cu_loc, work_tree, tree_type, new_split);
+        split_cost += search_cu(state, &new_cu_loc[1], work_tree, tree_type, new_split);
       }
       if (split_cost < cost) {
-        uvg_cu_loc_ctor(&new_cu_loc, x, y + half_cu, half_cu, half_cu);
-        split_cost += search_cu(state, &new_cu_loc, work_tree, tree_type, new_split);
+        split_cost += search_cu(state, &new_cu_loc[2], work_tree, tree_type, new_split);
       }
       if (split_cost < cost) {
-        uvg_cu_loc_ctor(&new_cu_loc, x + half_cu, y + half_cu, half_cu, half_cu);
-        split_cost += search_cu(state, &new_cu_loc, work_tree, tree_type, new_split);
+        split_cost += search_cu(state, &new_cu_loc[3], work_tree, tree_type, new_split);
       }
     } else {
       split_cost = INT_MAX;