1 files changed, 33 insertions, 34 deletions

diff --git a/third_party/aom/av1/encoder/context_tree.h b/third_party/aom/av1/encoder/context_tree.h
index 38052ba27..c05f48a7a 100644
--- a/third_party/aom/av1/encoder/context_tree.h
+++ b/third_party/aom/av1/encoder/context_tree.h
@@ -23,28 +23,29 @@ struct AV1_COMP;
 struct AV1Common;
 struct ThreadData;
 
+typedef enum {
+  // Search all the partition types in this plane.
+  SEARCH_FULL_PLANE = 0,
+  // Only search none_partition coding block.
+  NONE_PARTITION_PLANE = 1,
+  // Search all the partition types in this plane except split.
+  SEARCH_SAME_PLANE = 2,
+  // Skip search partition on this plane. Go split directly.
+  SPLIT_PLANE = 3,
+} CB_TREE_SEARCH;
+
 // Structure to hold snapshot of coding context during the mode picking process
 typedef struct {
-  MODE_INFO mic;
+  MB_MODE_INFO mic;
   MB_MODE_INFO_EXT mbmi_ext;
   uint8_t *color_index_map[2];
-#if CONFIG_MRC_TX
-  uint8_t *mrc_mask;
-#endif  // CONFIG_MRC_TX
-#if CONFIG_VAR_TX
-  uint8_t *blk_skip[MAX_MB_PLANE];
-#endif
+  uint8_t *blk_skip;
 
   tran_low_t *coeff[MAX_MB_PLANE];
   tran_low_t *qcoeff[MAX_MB_PLANE];
   tran_low_t *dqcoeff[MAX_MB_PLANE];
-#if CONFIG_PVQ
-  tran_low_t *pvq_ref_coeff[MAX_MB_PLANE];
-#endif
   uint16_t *eobs[MAX_MB_PLANE];
-#if CONFIG_LV_MAP
   uint8_t *txb_entropy_ctx[MAX_MB_PLANE];
-#endif
 
   int num_4x4_blk;
   int skip;
@@ -60,16 +61,27 @@ typedef struct {
   // scope of refactoring.
   int rate;
   int64_t dist;
+  int64_t rdcost;
+  int rd_mode_is_ready;  // Flag to indicate whether rd pick mode decision has
+                         // been made.
 
   // motion vector cache for adaptive motion search control in partition
   // search loop
-  MV pred_mv[TOTAL_REFS_PER_FRAME];
+  MV pred_mv[REF_FRAMES];
   InterpFilter pred_interp_filter;
-#if CONFIG_EXT_PARTITION_TYPES
   PARTITION_TYPE partition;
-#endif
 } PICK_MODE_CONTEXT;
 
+typedef struct {
+  int valid;
+  int split;
+  int skip;
+  int64_t rdcost;
+  int sub_block_split[4];
+  int sub_block_skip[4];
+  int64_t sub_block_rdcost[4];
+} PC_TREE_STATS;
+
 typedef struct PC_TREE {
   int index;
   PARTITION_TYPE partitioning;
@@ -77,34 +89,21 @@ typedef struct PC_TREE {
   PICK_MODE_CONTEXT none;
   PICK_MODE_CONTEXT horizontal[2];
   PICK_MODE_CONTEXT vertical[2];
-#if CONFIG_EXT_PARTITION_TYPES
   PICK_MODE_CONTEXT horizontala[3];
   PICK_MODE_CONTEXT horizontalb[3];
   PICK_MODE_CONTEXT verticala[3];
   PICK_MODE_CONTEXT verticalb[3];
   PICK_MODE_CONTEXT horizontal4[4];
   PICK_MODE_CONTEXT vertical4[4];
-#endif
-  // TODO(jingning): remove leaf_split[] when cb4x4 experiment flag is removed.
-  union {
-    struct PC_TREE *split[4];
-    PICK_MODE_CONTEXT *leaf_split[4];
-  };
-#if CONFIG_SUPERTX
-  PICK_MODE_CONTEXT horizontal_supertx;
-  PICK_MODE_CONTEXT vertical_supertx;
-  PICK_MODE_CONTEXT split_supertx;
-#if CONFIG_EXT_PARTITION_TYPES
-  PICK_MODE_CONTEXT horizontala_supertx;
-  PICK_MODE_CONTEXT horizontalb_supertx;
-  PICK_MODE_CONTEXT verticala_supertx;
-  PICK_MODE_CONTEXT verticalb_supertx;
-#endif
-#endif
+  CB_TREE_SEARCH cb_search_range;
+  struct PC_TREE *split[4];
+  PC_TREE_STATS pc_tree_stats;
 } PC_TREE;
 
 void av1_setup_pc_tree(struct AV1Common *cm, struct ThreadData *td);
-void av1_free_pc_tree(struct ThreadData *td);
+void av1_free_pc_tree(struct ThreadData *td, const int num_planes);
+void av1_copy_tree_context(PICK_MODE_CONTEXT *dst_ctx,
+                           PICK_MODE_CONTEXT *src_ctx);
 
 #ifdef __cplusplus
 }  // extern "C"