Refactor PC_TREE root handling.

Change-Id: Id8b16c1b18bd6f909e72aae3fd582dd3503c88c6
diff --git a/vp10/common/enums.h b/vp10/common/enums.h
index e144a45..86a7efc 100644
--- a/vp10/common/enums.h
+++ b/vp10/common/enums.h
@@ -20,7 +20,7 @@
 
 #undef MAX_SB_SIZE
 
-// Pixels per max superblock size
+// Max superblock size
 #if CONFIG_EXT_PARTITION
 # define MAX_SB_SIZE_LOG2 7
 #else
@@ -29,6 +29,9 @@
 #define MAX_SB_SIZE   (1 << MAX_SB_SIZE_LOG2)
 #define MAX_SB_SQUARE (MAX_SB_SIZE * MAX_SB_SIZE)
 
+// Min superblock size
+#define MIN_SB_SIZE_LOG2 6
+
 // Pixels per Mode Info (MI) unit
 #define MI_SIZE_LOG2  3
 #define MI_SIZE       (1 << MI_SIZE_LOG2)
@@ -37,6 +40,9 @@
 #define MAX_MIB_SIZE_LOG2 (MAX_SB_SIZE_LOG2 - MI_SIZE_LOG2)
 #define MAX_MIB_SIZE      (1 << MAX_MIB_SIZE_LOG2)
 
+// MI-units per min superblock
+#define MIN_MIB_SIZE_LOG2 (MIN_SB_SIZE_LOG2 - MI_SIZE_LOG2)
+
 // Mask to extract MI offset within max MIB
 #define MAX_MIB_MASK    (MAX_MIB_SIZE - 1)
 #define MAX_MIB_MASK_2  (MAX_MIB_SIZE * 2 - 1)
diff --git a/vp10/encoder/context_tree.c b/vp10/encoder/context_tree.c
index b7c8260..41155c9 100644
--- a/vp10/encoder/context_tree.c
+++ b/vp10/encoder/context_tree.c
@@ -244,8 +244,16 @@
     }
     ++square_index;
   }
-  td->pc_root = &td->pc_tree[tree_nodes - 1];
-  td->pc_root[0].none.best_mode_index = 2;
+
+  // Set up the root node for the largest superblock size
+  i = MAX_MIB_SIZE_LOG2 - MIN_MIB_SIZE_LOG2;
+  td->pc_root[i] = &td->pc_tree[tree_nodes - 1];
+  td->pc_root[i]->none.best_mode_index = 2;
+  // Set up the root nodes for the rest of the possible superblock sizes
+  while (--i >= 0) {
+    td->pc_root[i] = td->pc_root[i+1]->split[0];
+    td->pc_root[i]->none.best_mode_index = 2;
+  }
 }
 
 void vp10_free_pc_tree(ThreadData *td) {
diff --git a/vp10/encoder/encodeframe.c b/vp10/encoder/encodeframe.c
index 6aba475..88e9486 100644
--- a/vp10/encoder/encodeframe.c
+++ b/vp10/encoder/encodeframe.c
@@ -4235,6 +4235,7 @@
 
     const int idx_str = cm->mi_stride * mi_row + mi_col;
     MODE_INFO **mi = cm->mi_grid_visible + idx_str;
+    PC_TREE *const pc_root = td->pc_root[cm->mib_size_log2 - MIN_MIB_SIZE_LOG2];
 
     if (sf->adaptive_pred_interp_filter) {
       for (i = 0; i < leaf_nodes; ++i)
@@ -4249,7 +4250,7 @@
     }
 
     vp10_zero(x->pred_mv);
-    td->pc_root->index = 0;
+    pc_root->index = 0;
 
     if (seg->enabled) {
       const uint8_t *const map = seg->update_map ? cpi->segmentation_map
@@ -4269,7 +4270,7 @@
 #if CONFIG_SUPERTX
                        &dummy_rate_nocoef,
 #endif  // CONFIG_SUPERTX
-                       1, td->pc_root);
+                       1, pc_root);
     } else if (cpi->partition_search_skippable_frame) {
       BLOCK_SIZE bsize;
       set_offsets(cpi, tile_info, x, mi_row, mi_col, cm->sb_size);
@@ -4280,7 +4281,7 @@
 #if CONFIG_SUPERTX
                        &dummy_rate_nocoef,
 #endif  // CONFIG_SUPERTX
-                       1, td->pc_root);
+                       1, pc_root);
     } else if (sf->partition_search_type == VAR_BASED_PARTITION &&
                cm->frame_type != KEY_FRAME) {
       choose_partitioning(cpi, tile_info, x, mi_row, mi_col);
@@ -4289,7 +4290,7 @@
 #if CONFIG_SUPERTX
                        &dummy_rate_nocoef,
 #endif  // CONFIG_SUPERTX
-                       1, td->pc_root);
+                       1, pc_root);
     } else {
       // If required set upper and lower partition size limits
       if (sf->auto_min_max_partition_size) {
@@ -4303,9 +4304,7 @@
 #if CONFIG_SUPERTX
                         &dummy_rate_nocoef,
 #endif  // CONFIG_SUPERTX
-                        INT64_MAX,
-                        cm->sb_size == BLOCK_LARGEST ? td->pc_root
-                                                     : td->pc_root->split[0]);
+                        INT64_MAX, pc_root);
     }
   }
 #if CONFIG_ENTROPY
diff --git a/vp10/encoder/encoder.h b/vp10/encoder/encoder.h
index 2098378..bf7815f 100644
--- a/vp10/encoder/encoder.h
+++ b/vp10/encoder/encoder.h
@@ -266,7 +266,7 @@
 
   PICK_MODE_CONTEXT *leaf_tree;
   PC_TREE *pc_tree;
-  PC_TREE *pc_root;
+  PC_TREE *pc_root[MAX_MIB_SIZE_LOG2 - MIN_MIB_SIZE_LOG2 + 1];
 } ThreadData;
 
 struct EncWorkerData;
diff --git a/vp10/encoder/firstpass.c b/vp10/encoder/firstpass.c
index dd3e437..5936a24 100644
--- a/vp10/encoder/firstpass.c
+++ b/vp10/encoder/firstpass.c
@@ -491,7 +491,8 @@
   TileInfo tile;
   struct macroblock_plane *const p = x->plane;
   struct macroblockd_plane *const pd = xd->plane;
-  const PICK_MODE_CONTEXT *ctx = &cpi->td.pc_root->none;
+  const PICK_MODE_CONTEXT *ctx =
+      &cpi->td.pc_root[MAX_MIB_SIZE_LOG2 - MIN_MIB_SIZE_LOG2]->none;
   int i;
 
   int recon_yoffset, recon_uvoffset;