ext_partition_types: Pass the correct CDF length for partitions

Each CDF for partitioning square blocks is initialised from
an entry of default_partition_cdf in entropymode.c. These CDFs are of
different lengths, depending on which partition types are supported by
the block size.

For example, 8x8 blocks have a CDF with only 4 entries (PARTITION_NONE
through PARTITION_SPLIT). Blocks of a size that supports 1:4 and 4:1
partitions have 10 entries. Currently, that's only 32x32 blocks. All
other blocks have 8 entries.

Change-Id: Ie2126b6d41afc0efedcc5b5b37fc1d0427b9a9fa
diff --git a/av1/common/entropymode.c b/av1/common/entropymode.c
index 2f94f35..94b7f74 100644
--- a/av1/common/entropymode.c
+++ b/av1/common/entropymode.c
@@ -5160,6 +5160,18 @@
 int av1_switchable_interp_ind[SWITCHABLE_FILTERS];
 int av1_switchable_interp_inv[SWITCHABLE_FILTERS];
 
+#if CONFIG_EXT_PARTITION_TYPES
+int av1_num_partition_types[PARTITION_BLOCK_SIZES] = {
+  PARTITION_TYPES,          // 8x8: The 4 traditional partitions
+  EXT_PARTITION_TYPES - 2,  // 16x16: All but 4:1 and 1:4 partitions
+  EXT_PARTITION_TYPES,      // 32x32: All partitions
+  EXT_PARTITION_TYPES - 2,  // 64x64: All but 4:1 and 1:4 partitions
+#if CONFIG_EXT_PARTITION
+  EXT_PARTITION_TYPES - 2  // 128x128: All but 4:1 and 1:4 partitions
+#endif
+};
+#endif  // CONFIG_EXT_PARTITION_TYPES
+
 #if CONFIG_DUAL_FILTER && USE_EXTRA_FILTER
 const aom_tree_index av1_switchable_interp_tree[TREE_SIZE(SWITCHABLE_FILTERS)] =
     {
diff --git a/av1/common/entropymode.h b/av1/common/entropymode.h
index 0a07087..8fb8089 100644
--- a/av1/common/entropymode.h
+++ b/av1/common/entropymode.h
@@ -585,6 +585,10 @@
 extern int av1_switchable_interp_ind[SWITCHABLE_FILTERS];
 extern int av1_switchable_interp_inv[SWITCHABLE_FILTERS];
 
+#if CONFIG_EXT_PARTITION_TYPES
+extern int av1_num_partition_types[PARTITION_BLOCK_SIZES];
+#endif
+
 void av1_setup_past_independence(struct AV1Common *cm);
 
 void av1_adapt_intra_frame_probs(struct AV1Common *cm);
diff --git a/av1/decoder/decodeframe.c b/av1/decoder/decodeframe.c
index 5740f53..07e4ea4 100644
--- a/av1/decoder/decodeframe.c
+++ b/av1/decoder/decodeframe.c
@@ -2291,24 +2291,23 @@
 
   aom_cdf_prob *partition_cdf = (ctx >= 0) ? ec_ctx->partition_cdf[ctx] : NULL;
 
-  if (has_rows && has_cols)
+  if (has_rows && has_cols) {
 #if CONFIG_EXT_PARTITION_TYPES
-    if (bsize <= BLOCK_8X8)
-      p = (PARTITION_TYPE)aom_read_symbol(r, partition_cdf, PARTITION_TYPES,
-                                          ACCT_STR);
-    else
-      p = (PARTITION_TYPE)aom_read_symbol(r, partition_cdf, EXT_PARTITION_TYPES,
-                                          ACCT_STR);
+    const int bsl =
+        mi_width_log2_lookup[bsize] - mi_width_log2_lookup[BLOCK_8X8];
+    p = (PARTITION_TYPE)aom_read_symbol(r, partition_cdf,
+                                        av1_num_partition_types[bsl], ACCT_STR);
 #else
     p = (PARTITION_TYPE)aom_read_symbol(r, partition_cdf, PARTITION_TYPES,
                                         ACCT_STR);
 #endif  // CONFIG_EXT_PARTITION_TYPES
-  else if (!has_rows && has_cols)
+  } else if (!has_rows && has_cols) {
     p = aom_read(r, probs[1], ACCT_STR) ? PARTITION_SPLIT : PARTITION_HORZ;
-  else if (has_rows && !has_cols)
+  } else if (has_rows && !has_cols) {
     p = aom_read(r, probs[2], ACCT_STR) ? PARTITION_SPLIT : PARTITION_VERT;
-  else
+  } else {
     p = PARTITION_SPLIT;
+  }
 
   if (counts) ++counts->partition[ctx][p];
 
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index 11acede..26b72d0 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -2847,10 +2847,10 @@
 
   if (has_rows && has_cols) {
 #if CONFIG_EXT_PARTITION_TYPES
-    if (bsize <= BLOCK_8X8)
-      aom_write_symbol(w, p, ec_ctx->partition_cdf[ctx], PARTITION_TYPES);
-    else
-      aom_write_symbol(w, p, ec_ctx->partition_cdf[ctx], EXT_PARTITION_TYPES);
+    const int bsl =
+        mi_width_log2_lookup[bsize] - mi_width_log2_lookup[BLOCK_8X8];
+    aom_write_symbol(w, p, ec_ctx->partition_cdf[ctx],
+                     av1_num_partition_types[bsl]);
 #else
     aom_write_symbol(w, p, ec_ctx->partition_cdf[ctx], PARTITION_TYPES);
 #endif  // CONFIG_EXT_PARTITION_TYPES