Simplify get_partition() implementation

This function is given a bsize and an mi array and has to figure out
what partition to use to divide the given bsize in the direction of
the sizes it finds in the mi array. (Since each block size can be
reached by only one sequence of partitions, this can be done
unambiguously)

The previous version was correct, working by looking up entries in the
partition_lookup array. Unfortunately, that lookup isn't quite enough
when CONFIG_EXT_PARTITION_TYPES is true, so it then had to do some
slightly confusing computations to fix things up after the fact.

The new version should be more self-explanatory and doesn't work by
looking things up in a magic array. It looks up the width and height
corresponding to bsize and compares them with the width and height
corresponding to the sb_type at mi_row,mi_col in the mi array. When
CONFIG_EXT_PARTITION_TYPES is false, this is all you need, and the
four corresponding cases can be found by a lookup in an array of 4
elements.

With extended partition types and a sufficiently large block, you have
to do a bit more searching. For example, if bsize is BLOCK_16X16 and
the subsize is BLOCK_8X8, the partition might be PARTITION_SPLIT, but
it might be one of PARTITION_HORZ_A or PARTITION_VERT_A instead. The
new code adds some comments to explain what's going on.

A nice side-effect of rewriting get_partition in this way is that it
lets us completely dispense with the partition_lookup array.

The patch also fixes comments for the A/B extended partitions in
enums.h, which were slightly backwards (a "horizontal split means two
blocks vertically above one another)

Change-Id: I4b48189103aa63e1859f25a15d7690d53ca7baf5
diff --git a/av1/common/common_data.h b/av1/common/common_data.h
index ee8d3bf..1e5e184 100644
--- a/av1/common/common_data.h
+++ b/av1/common/common_data.h
@@ -174,176 +174,6 @@
 };
 
 /* clang-format off */
-static const PARTITION_TYPE
-  partition_lookup[MAX_SB_SIZE_LOG2 - 1][BLOCK_SIZES_ALL] = {
-  {     // 4X4 ->
-#if CONFIG_CHROMA_2X2 || CONFIG_CHROMA_SUB8X8
-    // 2X2,            2X4,               4X2,
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#endif
-    //                                    4X4
-                                          PARTITION_NONE,
-    // 4X8,            8X4,               8X8
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-    // 8X16,           16X8,              16X16
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-    // 16X32,          32X16,             32X32
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-    // 32X64,          64X32,             64X64
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#if CONFIG_EXT_PARTITION
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#endif  // CONFIG_EXT_PARTITION
-    // 4x16,           16x4,              8x32
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-    // 32x8            16x64              64x16
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-  }, {  // 8X8 ->
-#if CONFIG_CHROMA_2X2 || CONFIG_CHROMA_SUB8X8
-    // 2X2,            2X4,               4X2,
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#endif
-    //                                    4X4
-                                          PARTITION_SPLIT,
-    // 4X8,            8X4,               8X8
-    PARTITION_VERT,    PARTITION_HORZ,    PARTITION_NONE,
-    // 8X16,           16X8,              16X16
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-    // 16X32,          32X16,             32X32
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-    // 32X64,          64X32,             64X64
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#if CONFIG_EXT_PARTITION
-    // 64x128,         128x64,            128x128
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#endif  // CONFIG_EXT_PARTITION
-    // 4x16,           16x4,              8x32
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-    // 32x8            16x64              64x16
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-  }, {  // 16X16 ->
-#if CONFIG_CHROMA_2X2 || CONFIG_CHROMA_SUB8X8
-    // 2X2,            2X4,               4X2,
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#endif
-    //                                    4X4
-                                          PARTITION_SPLIT,
-    // 4X8,            8X4,               8X8
-    PARTITION_SPLIT,   PARTITION_SPLIT,   PARTITION_SPLIT,
-    // 8X16,           16X8,              16X16
-    PARTITION_VERT,    PARTITION_HORZ,    PARTITION_NONE,
-    // 16X32,          32X16,             32X32
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-    // 32X64,          64X32,             64X64
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#if CONFIG_EXT_PARTITION
-    // 64x128,         128x64,            128x128
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#endif  // CONFIG_EXT_PARTITION
-#if CONFIG_EXT_PARTITION_TYPES
-    // 4x16,           16x4,              8x32
-    PARTITION_VERT_4,  PARTITION_HORZ_4,  PARTITION_INVALID,
-    // 32x8            16x64              64x16
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#else
-    // 4x16,           16x4,              8x32
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-    // 32x8            16x64              64x16
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#endif
-  }, {  // 32X32 ->
-#if CONFIG_CHROMA_2X2 || CONFIG_CHROMA_SUB8X8
-    // 2X2,            2X4,               4X2,
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#endif
-    //                                    4X4
-                                          PARTITION_SPLIT,
-    // 4X8,            8X4,               8X8
-    PARTITION_SPLIT,   PARTITION_SPLIT,   PARTITION_SPLIT,
-    // 8X16,           16X8,              16X16
-    PARTITION_SPLIT,   PARTITION_SPLIT,   PARTITION_SPLIT,
-    // 16X32,          32X16,             32X32
-    PARTITION_VERT,    PARTITION_HORZ,    PARTITION_NONE,
-    // 32X64,          64X32,             64X64
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#if CONFIG_EXT_PARTITION
-    // 64x128,         128x64,            128x128
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#endif  // CONFIG_EXT_PARTITION
-#if CONFIG_EXT_PARTITION_TYPES
-    // 4x16,           16x4,              8x32
-    PARTITION_SPLIT,   PARTITION_SPLIT,   PARTITION_VERT_4,
-    // 32x8            16x64              64x16
-    PARTITION_HORZ_4,  PARTITION_INVALID, PARTITION_INVALID,
-#else
-    // 4x16,           16x4,              8x32
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-    // 32x8            16x64              64x16
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#endif
-  }, {  // 64X64 ->
-#if CONFIG_CHROMA_2X2 || CONFIG_CHROMA_SUB8X8
-    // 2X2,            2X4,               4X2,
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#endif
-    //                                    4X4
-                                          PARTITION_SPLIT,
-    // 4X8,            8X4,               8X8
-    PARTITION_SPLIT,   PARTITION_SPLIT,   PARTITION_SPLIT,
-    // 8X16,           16X8,              16X16
-    PARTITION_SPLIT,   PARTITION_SPLIT,   PARTITION_SPLIT,
-    // 16X32,          32X16,             32X32
-    PARTITION_SPLIT,   PARTITION_SPLIT,   PARTITION_SPLIT,
-    // 32X64,          64X32,             64X64
-    PARTITION_VERT,    PARTITION_HORZ,    PARTITION_NONE,
-#if CONFIG_EXT_PARTITION
-    // 64x128,         128x64,            128x128
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#endif  // CONFIG_EXT_PARTITION
-#if CONFIG_EXT_PARTITION_TYPES
-    // 4x16,           16x4,              8x32
-    PARTITION_SPLIT,   PARTITION_SPLIT,   PARTITION_SPLIT,
-    // 32x8            16x64              64x16
-    PARTITION_SPLIT,   PARTITION_VERT_4,  PARTITION_HORZ_4
-#else
-    // 4x16,           16x4,              8x32
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-    // 32x8            16x64              64x16
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#endif  // CONFIG_EXT_PARTITION_TYPES
-#if CONFIG_EXT_PARTITION
-  }, {  // 128x128 ->
-#if CONFIG_CHROMA_2X2 || CONFIG_CHROMA_SUB8X8
-    // 2X2,            2X4,               4X2,
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-#endif
-    //                                    4X4
-                                          PARTITION_SPLIT,
-    // 4X8,            8X4,               8X8
-    PARTITION_SPLIT,   PARTITION_SPLIT,   PARTITION_SPLIT,
-    // 8X16,           16X8,              16X16
-    PARTITION_SPLIT,   PARTITION_SPLIT,   PARTITION_SPLIT,
-    // 16X32,          32X16,             32X32
-    PARTITION_SPLIT,   PARTITION_SPLIT,   PARTITION_SPLIT,
-    // 32X64,          64X32,             64X64
-    PARTITION_SPLIT,   PARTITION_SPLIT,   PARTITION_SPLIT,
-    // 64x128,         128x64,            128x128
-    PARTITION_VERT,    PARTITION_HORZ,    PARTITION_NONE,
-#if CONFIG_EXT_PARTITION_TYPES
-    // 4x16,           16x4,              8x32
-    PARTITION_SPLIT,   PARTITION_SPLIT,   PARTITION_SPLIT,
-    // 32x8            16x64              64x16
-    PARTITION_SPLIT,   PARTITION_SPLIT,   PARTITION_SPLIT
-#else
-    // 4x16,           16x4,              8x32
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID,
-    // 32x8            16x64              64x16
-    PARTITION_INVALID, PARTITION_INVALID, PARTITION_INVALID
-#endif  // CONFIG_EXT_PARTITION_TYPES
-#endif  // CONFIG_EXT_PARTITION
-  }
-};
-
 #if CONFIG_EXT_PARTITION_TYPES
 static const BLOCK_SIZE subsize_lookup[EXT_PARTITION_TYPES][BLOCK_SIZES_ALL] =
 #else
diff --git a/av1/common/enums.h b/av1/common/enums.h
index 31c54fb..fda6a86 100644
--- a/av1/common/enums.h
+++ b/av1/common/enums.h
@@ -137,10 +137,10 @@
   PARTITION_VERT,
   PARTITION_SPLIT,
 #if CONFIG_EXT_PARTITION_TYPES
-  PARTITION_HORZ_A,  // HORZ split and the left partition is split again
-  PARTITION_HORZ_B,  // HORZ split and the right partition is split again
-  PARTITION_VERT_A,  // VERT split and the top partition is split again
-  PARTITION_VERT_B,  // VERT split and the bottom partition is split again
+  PARTITION_HORZ_A,  // HORZ split and the top partition is split again
+  PARTITION_HORZ_B,  // HORZ split and the bottom partition is split again
+  PARTITION_VERT_A,  // VERT split and the left partition is split again
+  PARTITION_VERT_B,  // VERT split and the right partition is split again
   PARTITION_HORZ_4,  // 4:1 horizontal partition
   PARTITION_VERT_4,  // 4:1 vertical partition
   EXT_PARTITION_TYPES,
diff --git a/av1/common/onyxc_int.h b/av1/common/onyxc_int.h
index 4dc7a3e..ebf7ea4 100644
--- a/av1/common/onyxc_int.h
+++ b/av1/common/onyxc_int.h
@@ -1161,49 +1161,81 @@
 }
 #endif
 
+// Compute the next partition in the direction of the sb_type stored in the mi
+// array, starting with bsize.
 static INLINE PARTITION_TYPE get_partition(const AV1_COMMON *const cm,
                                            int mi_row, int mi_col,
                                            BLOCK_SIZE bsize) {
-  if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols) {
-    return PARTITION_INVALID;
-  } else {
-    const int offset = mi_row * cm->mi_stride + mi_col;
-    MODE_INFO **mi = cm->mi_grid_visible + offset;
-    const MB_MODE_INFO *const mbmi = &mi[0]->mbmi;
-    const int bsl = b_width_log2_lookup[bsize];
-    const PARTITION_TYPE partition = partition_lookup[bsl][mbmi->sb_type];
-#if !CONFIG_EXT_PARTITION_TYPES
-    return partition;
-#else
-    const int hbs = mi_size_wide[bsize] / 2;
+  if (mi_row >= cm->mi_rows || mi_col >= cm->mi_cols) return PARTITION_INVALID;
 
-    assert(cm->mi_grid_visible[offset] == &cm->mi[offset]);
+  const int offset = mi_row * cm->mi_stride + mi_col;
+  MODE_INFO **mi = cm->mi_grid_visible + offset;
+  const BLOCK_SIZE subsize = mi[0]->mbmi.sb_type;
 
-    if (partition == PARTITION_HORZ_4 || partition == PARTITION_VERT_4)
-      return partition;
+  if (subsize == bsize) return PARTITION_NONE;
 
-    if (partition != PARTITION_NONE && bsize > BLOCK_8X8 &&
-        mi_row + hbs < cm->mi_rows && mi_col + hbs < cm->mi_cols) {
-      const BLOCK_SIZE h = get_subsize(bsize, PARTITION_HORZ_A);
-      const BLOCK_SIZE v = get_subsize(bsize, PARTITION_VERT_A);
-      const MB_MODE_INFO *const mbmi_right = &mi[hbs]->mbmi;
-      const MB_MODE_INFO *const mbmi_below = &mi[hbs * cm->mi_stride]->mbmi;
-      if (mbmi->sb_type == h) {
-        return mbmi_below->sb_type == h ? PARTITION_HORZ : PARTITION_HORZ_B;
-      } else if (mbmi->sb_type == v) {
-        return mbmi_right->sb_type == v ? PARTITION_VERT : PARTITION_VERT_B;
-      } else if (mbmi_below->sb_type == h) {
-        return PARTITION_HORZ_A;
-      } else if (mbmi_right->sb_type == v) {
-        return PARTITION_VERT_A;
-      } else {
-        return PARTITION_SPLIT;
-      }
+  const int bhigh = mi_size_high[bsize];
+  const int bwide = mi_size_wide[bsize];
+  const int sshigh = mi_size_high[subsize];
+  const int sswide = mi_size_wide[subsize];
+
+#if CONFIG_EXT_PARTITION_TYPES
+  if (bsize > BLOCK_8X8 && mi_row + bwide / 2 < cm->mi_rows &&
+      mi_col + bhigh / 2 < cm->mi_cols) {
+    // In this case, the block might be using an extended partition
+    // type.
+    const MB_MODE_INFO *const mbmi_right = &mi[bwide / 2]->mbmi;
+    const MB_MODE_INFO *const mbmi_below = &mi[bhigh / 2 * cm->mi_stride]->mbmi;
+
+    if (sswide == bwide) {
+      // Smaller height but same width. Is PARTITION_HORZ_4, PARTITION_HORZ or
+      // PARTITION_HORZ_B. To distinguish the latter two, check if the lower
+      // half was split.
+      if (sshigh * 4 == bhigh) return PARTITION_HORZ_4;
+      assert(sshigh * 2 == bhigh);
+
+      if (mbmi_below->sb_type == subsize)
+        return PARTITION_HORZ;
+      else
+        return PARTITION_HORZ_B;
+    } else if (sshigh == bhigh) {
+      // Smaller width but same height. Is PARTITION_VERT_4, PARTITION_VERT or
+      // PARTITION_VERT_B. To distinguish the latter two, check if the right
+      // half was split.
+      if (sswide * 4 == bwide) return PARTITION_VERT_4;
+      assert(sswide * 2 == bhigh);
+
+      if (mbmi_right->sb_type == subsize)
+        return PARTITION_VERT;
+      else
+        return PARTITION_VERT_B;
+    } else {
+      // Smaller width and smaller height. Might be PARTITION_SPLIT or could be
+      // PARTITION_HORZ_A or PARTITION_VERT_A. If subsize isn't halved in both
+      // dimensions, we immediately know this is a split (which will recurse to
+      // get to subsize). Otherwise look down and to the right. With
+      // PARTITION_VERT_A, the right block will have height bhigh; with
+      // PARTITION_HORZ_A, the lower block with have width bwide. Otherwise
+      // it's PARTITION_SPLIT.
+      if (sswide * 2 != bwide || sshigh * 2 != bhigh) return PARTITION_SPLIT;
+
+      if (mi_size_wide[mbmi_below->sb_type] == bwide) return PARTITION_HORZ_A;
+      if (mi_size_high[mbmi_right->sb_type] == bhigh) return PARTITION_VERT_A;
+
+      return PARTITION_SPLIT;
     }
-
-    return partition;
-#endif  // !CONFIG_EXT_PARTITION_TYPES
   }
+#endif
+  const int vert_split = sswide < bwide;
+  const int horz_split = sshigh < bhigh;
+  const int split_idx = (vert_split << 1) | horz_split;
+  assert(split_idx != 0);
+
+  static const PARTITION_TYPE base_partitions[4] = {
+    PARTITION_INVALID, PARTITION_HORZ, PARTITION_VERT, PARTITION_SPLIT
+  };
+
+  return base_partitions[split_idx];
 }
 
 static INLINE void set_sb_size(AV1_COMMON *const cm, BLOCK_SIZE sb_size) {