From 4e18d403d8c0e20f9aa84dfe870a3749d672e2af Mon Sep 17 00:00:00 2001
From: Alex Converse <aconverse@google.com>
Date: Tue, 14 Mar 2017 15:36:38 -0700
Subject: [PATCH] ext_partition_types: Make update tree shape match

Update the first PARTITION_PLOFFSET (4) contexts with the four classic
partitions. The extended partitions are only codable above 8x8, but
there are PARTITION_PLOFFSET (4) contexts for dropping below 8x8.

Change-Id: Ib3291dded6dc24103222e8f470504c20e29adb88
---
 av1/common/entropymode.c  |  9 +++++----
 av1/decoder/decodeframe.c |  7 ++++---
 av1/encoder/bitstream.c   | 10 +++++-----
 av1/encoder/rd.c          |  7 ++++---
 4 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/av1/common/entropymode.c b/av1/common/entropymode.c
index b6e4546fa..ec3dc28cf 100644
--- a/av1/common/entropymode.c
+++ b/av1/common/entropymode.c
@@ -2528,13 +2528,14 @@ void av1_adapt_intra_frame_probs(AV1_COMMON *cm) {
                          counts->uv_mode[i], fc->uv_mode_prob[i]);
 
 #if CONFIG_EXT_PARTITION_TYPES
-  aom_tree_merge_probs(av1_partition_tree, pre_fc->partition_prob[0],
-                       counts->partition[0], fc->partition_prob[0]);
-  for (i = 1; i < PARTITION_CONTEXTS_PRIMARY; i++)
+  for (i = 0; i < PARTITION_PLOFFSET; ++i)
+    aom_tree_merge_probs(av1_partition_tree, pre_fc->partition_prob[i],
+                         counts->partition[i], fc->partition_prob[i]);
+  for (; i < PARTITION_CONTEXTS_PRIMARY; ++i)
     aom_tree_merge_probs(av1_ext_partition_tree, pre_fc->partition_prob[i],
                          counts->partition[i], fc->partition_prob[i]);
 #else
-  for (i = 0; i < PARTITION_CONTEXTS_PRIMARY; i++) {
+  for (i = 0; i < PARTITION_CONTEXTS_PRIMARY; ++i) {
     aom_tree_merge_probs(av1_partition_tree, pre_fc->partition_prob[i],
                          counts->partition[i], fc->partition_prob[i]);
   }
diff --git a/av1/decoder/decodeframe.c b/av1/decoder/decodeframe.c
index b6bc87421..96afcb514 100644
--- a/av1/decoder/decodeframe.c
+++ b/av1/decoder/decodeframe.c
@@ -4611,9 +4611,10 @@ static int read_compressed_header(AV1Decoder *pbi, const uint8_t *data,
   }
 
 #if CONFIG_EXT_PARTITION_TYPES
-  for (i = 0; i < PARTITION_TYPES - 1; ++i)
-    av1_diff_update_prob(&r, &fc->partition_prob[0][i], ACCT_STR);
-  for (j = 1; j < PARTITION_CONTEXTS_PRIMARY; ++j)
+  for (j = 0; j < PARTITION_PLOFFSET; ++j)
+    for (i = 0; i < PARTITION_TYPES - 1; ++i)
+      av1_diff_update_prob(&r, &fc->partition_prob[j][i], ACCT_STR);
+  for (; j < PARTITION_CONTEXTS_PRIMARY; ++j)
     for (i = 0; i < EXT_PARTITION_TYPES - 1; ++i)
       av1_diff_update_prob(&r, &fc->partition_prob[j][i], ACCT_STR);
 #else
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index 12e8a818d..b7b82083e 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -4666,17 +4666,17 @@ static uint32_t write_compressed_header(AV1_COMP *cpi, uint8_t *data) {
   }
 
 #if CONFIG_EXT_PARTITION_TYPES
-  prob_diff_update(av1_partition_tree, fc->partition_prob[0],
-                   counts->partition[0], PARTITION_TYPES, probwt, header_bc);
-  for (i = 1; i < PARTITION_CONTEXTS_PRIMARY; ++i)
+  for (i = 0; i < PARTITION_PLOFFSET; ++i)
+    prob_diff_update(av1_partition_tree, fc->partition_prob[i],
+                     counts->partition[i], PARTITION_TYPES, probwt, header_bc);
+  for (; i < PARTITION_CONTEXTS_PRIMARY; ++i)
     prob_diff_update(av1_ext_partition_tree, fc->partition_prob[i],
                      counts->partition[i], EXT_PARTITION_TYPES, probwt,
                      header_bc);
 #else
-  for (i = 0; i < PARTITION_CONTEXTS_PRIMARY; ++i) {
+  for (i = 0; i < PARTITION_CONTEXTS_PRIMARY; ++i)
     prob_diff_update(av1_partition_tree, fc->partition_prob[i],
                      counts->partition[i], PARTITION_TYPES, probwt, header_bc);
-  }
 #endif  // CONFIG_EXT_PARTITION_TYPES
 #if CONFIG_UNPOISON_PARTITION_CTX
   for (; i < PARTITION_CONTEXTS_PRIMARY + PARTITION_BLOCK_SIZES; ++i) {
diff --git a/av1/encoder/rd.c b/av1/encoder/rd.c
index 76d18ea7c..d6561fc11 100644
--- a/av1/encoder/rd.c
+++ b/av1/encoder/rd.c
@@ -405,9 +405,10 @@ void av1_initialize_rd_consts(AV1_COMP *cpi) {
       cpi->partition_cost[0][PARTITION_SPLIT] = 0;
 #endif
 #if CONFIG_EXT_PARTITION_TYPES
-      av1_cost_tokens(cpi->partition_cost[CONFIG_UNPOISON_PARTITION_CTX],
-                      cm->fc->partition_prob[0], av1_partition_tree);
-      for (i = 1; i < PARTITION_CONTEXTS_PRIMARY; ++i)
+      for (i = 0; i < PARTITION_PLOFFSET; ++i)
+        av1_cost_tokens(cpi->partition_cost[CONFIG_UNPOISON_PARTITION_CTX + i],
+                        cm->fc->partition_prob[i], av1_partition_tree);
+      for (; i < PARTITION_CONTEXTS_PRIMARY; ++i)
         av1_cost_tokens(cpi->partition_cost[CONFIG_UNPOISON_PARTITION_CTX + i],
                         cm->fc->partition_prob[i], av1_ext_partition_tree);
 #else
-- 
GitLab