Adding {read, write}_partition() instead of check_bsize_coverage().

Making partition read/write logic more clear.

Change-Id: I1981e90327257d37095567c62d72a103cda1da33

vp9/common/vp9_onyxc_int.h

@@ -260,24 +260,6 @@ static INLINE void set_skip_context(
   }
 }
 
-// return the node index in the prob tree for binary coding
-static int check_bsize_coverage(int bs, int mi_rows, int mi_cols,
-                                int mi_row, int mi_col) {
-  const int r = (mi_row + bs < mi_rows);
-  const int c = (mi_col + bs < mi_cols);
-
-  if (r && c)
-    return 0;
-
-  if (c && !r)
-    return 1;  // only allow horizontal/split partition types
-
-  if (r && !c)
-    return 2;  // only allow vertical/split partition types
-
-  return -1;
-}
-
 static void set_mi_row_col(MACROBLOCKD *xd, const TileInfo *const tile,
                            int mi_row, int bh,
                            int mi_col, int bw,

vp9/decoder/vp9_decodframe.c

@@ -422,6 +422,23 @@ static void decode_modes_b(VP9_COMMON *const cm, MACROBLOCKD *const xd,
   xd->corrupted |= vp9_reader_has_error(r);
 }
 
+static PARTITION_TYPE read_partition(int hbs, int mi_rows, int mi_cols,
+                                     int mi_row, int mi_col,
+                                     vp9_prob probs[PARTITION_TYPES - 1],
+                                     vp9_reader *r) {
+  const int has_rows = (mi_row + hbs) < mi_rows;
+  const int has_cols = (mi_col + hbs) < mi_cols;
+
+  if (has_rows && has_cols)
+    return treed_read(r, vp9_partition_tree, probs);
+  else if (!has_rows && has_cols)
+    return vp9_read(r, probs[1]) ? PARTITION_SPLIT : PARTITION_HORZ;
+  else if (has_rows && !has_cols)
+    return vp9_read(r, probs[2]) ? PARTITION_SPLIT : PARTITION_VERT;
+  else
+    return PARTITION_SPLIT;
+}
+
 static void decode_modes_sb(VP9_COMMON *const cm, MACROBLOCKD *const xd,
                             const TileInfo *const tile,
                             int mi_row, int mi_col,
@@ -437,23 +454,14 @@ static void decode_modes_sb(VP9_COMMON *const cm, MACROBLOCKD *const xd,
     if (index > 0)
       return;
   } else {
-    int pl;
-    const int idx = check_bsize_coverage(hbs, cm->mi_rows, cm->mi_cols,
-                                         mi_row, mi_col);
-    pl = partition_plane_context(xd->above_seg_context, xd->left_seg_context,
-                                 mi_row, mi_col, bsize);
-
-    if (idx == 0)
-      partition = treed_read(r, vp9_partition_tree,
-                             cm->fc.partition_prob[cm->frame_type][pl]);
-    else if (idx > 0 &&
-        !vp9_read(r, cm->fc.partition_prob[cm->frame_type][pl][idx]))
-      partition = (idx == 1) ? PARTITION_HORZ : PARTITION_VERT;
-    else
-      partition = PARTITION_SPLIT;
+    const int ctx = partition_plane_context(xd->above_seg_context,
+                                            xd->left_seg_context,
+                                            mi_row, mi_col, bsize);
+    partition = read_partition(hbs, cm->mi_rows, cm->mi_cols, mi_row, mi_col,
+                               cm->fc.partition_prob[cm->frame_type][ctx], r);
 
     if (!cm->frame_parallel_decoding_mode)
-      ++cm->counts.partition[pl][partition];
+      ++cm->counts.partition[ctx][partition];
   }
 
   subsize = get_subsize(bsize, partition);

vp9/encoder/vp9_bitstream.c

@@ -595,6 +595,28 @@ static void write_modes_b(VP9_COMP *cpi, const TileInfo *const tile,
   pack_mb_tokens(bc, tok, tok_end);
 }
 
+static void write_partition(PARTITION_TYPE partition,
+                            int hbs, int mi_rows, int mi_cols,
+                            int mi_row, int mi_col,
+                            vp9_prob probs[PARTITION_TYPES - 1],
+                            vp9_writer *w) {
+  const int has_rows = (mi_row + hbs) < mi_rows;
+  const int has_cols = (mi_col + hbs) < mi_cols;
+
+  if (has_rows && has_cols) {
+    write_token(w, vp9_partition_tree, probs,
+                &vp9_partition_encodings[partition]);
+  } else if (!has_rows && has_cols) {
+    assert(partition == PARTITION_SPLIT || partition == PARTITION_HORZ);
+    vp9_write(w, partition == PARTITION_SPLIT, probs[1]);
+  } else if (has_rows && !has_cols) {
+    assert(partition == PARTITION_SPLIT || partition == PARTITION_VERT);
+    vp9_write(w, partition == PARTITION_SPLIT, probs[2]);
+  } else {
+    assert(partition == PARTITION_SPLIT);
+  }
+}
+
 static void write_modes_sb(VP9_COMP *cpi, const TileInfo *const tile,
                            MODE_INFO **mi_8x8, vp9_writer *bc,
                            TOKENEXTRA **tok, TOKENEXTRA *tok_end,
@@ -618,19 +640,11 @@ static void write_modes_sb(VP9_COMP *cpi, const TileInfo *const tile,
     if (index > 0)
       return;
   } else {
-    int pl;
-    const int idx = check_bsize_coverage(bs, cm->mi_rows, cm->mi_cols,
-                                         mi_row, mi_col);
-    pl = partition_plane_context(cpi->above_seg_context, cpi->left_seg_context,
-                                 mi_row, mi_col, bsize);
-    // encode the partition information
-    if (idx == 0)
-      write_token(bc, vp9_partition_tree,
-                  cm->fc.partition_prob[cm->frame_type][pl],
-                  vp9_partition_encodings + partition);
-    else if (idx > 0)
-      vp9_write(bc, partition == PARTITION_SPLIT,
-                cm->fc.partition_prob[cm->frame_type][pl][idx]);
+    const int ctx = partition_plane_context(cpi->above_seg_context,
+                                            cpi->left_seg_context,
+                                            mi_row, mi_col, bsize);
+    write_partition(partition, bs, cm->mi_rows, cm->mi_cols, mi_row, mi_col,
+                    cm->fc.partition_prob[cm->frame_type][ctx], bc);
   }
 
   subsize = get_subsize(bsize, partition);