Merge "General cleanup in segmentation-related code."

6 files changed, 89 insertions, 122 deletions

diff --git a/vp9/common/vp9_pred_common.c b/vp9/common/vp9_pred_common.c
index d921a1197..3adfaebdb 100644
--- a/vp9/common/vp9_pred_common.c
+++ b/vp9/common/vp9_pred_common.c
@@ -19,8 +19,7 @@
 // TBD prediction functions for various bitstream signals
 
 // Returns a context number for the given MB prediction signal
-unsigned char vp9_get_pred_context(const VP9_COMMON *const cm,
-                                   const MACROBLOCKD *const xd,
+unsigned char vp9_get_pred_context(const VP9_COMMON *cm, const MACROBLOCKD *xd,
                                    PRED_ID pred_id) {
   int pred_context;
   const MODE_INFO *const mi = xd->mode_info_context;
@@ -389,9 +388,8 @@ unsigned char vp9_get_pred_context(const VP9_COMMON *const cm,
 
 // This function returns a context probability for coding a given
 // prediction signal
-vp9_prob vp9_get_pred_prob(const VP9_COMMON *const cm,
-                          const MACROBLOCKD *const xd,
-                          PRED_ID pred_id) {
+vp9_prob vp9_get_pred_prob(const VP9_COMMON *cm, const MACROBLOCKD *xd,
+                           PRED_ID pred_id) {
   const int pred_context = vp9_get_pred_context(cm, xd, pred_id);
 
   switch (pred_id) {
@@ -417,8 +415,7 @@ vp9_prob vp9_get_pred_prob(const VP9_COMMON *const cm,
 
 // This function returns a context probability ptr for coding a given
 // prediction signal
-const vp9_prob *vp9_get_pred_probs(const VP9_COMMON *const cm,
-                                   const MACROBLOCKD *const xd,
+const vp9_prob *vp9_get_pred_probs(const VP9_COMMON *cm, const MACROBLOCKD * xd,
                                    PRED_ID pred_id) {
   const MODE_INFO *const mi = xd->mode_info_context;
   const int pred_context = vp9_get_pred_context(cm, xd, pred_id);
@@ -458,8 +455,7 @@ unsigned char vp9_get_pred_flag(const MACROBLOCKD *const xd,
 
 // This function sets the status of the given prediction signal.
 // I.e. is the predicted value for the given signal correct.
-void vp9_set_pred_flag(MACROBLOCKD *const xd,
-                       PRED_ID pred_id,
+void vp9_set_pred_flag(MACROBLOCKD *xd, PRED_ID pred_id,
                        unsigned char pred_flag) {
   const int mis = xd->mode_info_stride;
   BLOCK_SIZE_TYPE bsize = xd->mode_info_context->mbmi.sb_type;
@@ -473,19 +469,15 @@ void vp9_set_pred_flag(MACROBLOCKD *const xd,
 
   switch (pred_id) {
     case PRED_SEG_ID:
-      for (y = 0; y < y_mis; y++) {
-        for (x = 0; x < x_mis; x++) {
+      for (y = 0; y < y_mis; y++)
+        for (x = 0; x < x_mis; x++)
           xd->mode_info_context[y * mis + x].mbmi.seg_id_predicted = pred_flag;
-        }
-      }
       break;
 
     case PRED_MBSKIP:
-      for (y = 0; y < y_mis; y++) {
-        for (x = 0; x < x_mis; x++) {
+      for (y = 0; y < y_mis; y++)
+        for (x = 0; x < x_mis; x++)
           xd->mode_info_context[y * mis + x].mbmi.mb_skip_coeff = pred_flag;
-        }
-      }
       break;
 
     default:
@@ -495,26 +487,20 @@ void vp9_set_pred_flag(MACROBLOCKD *const xd,
   }
 }
 
-
-// The following contain the guts of the prediction code used to
-// peredict various bitstream signals.
-
-// Macroblock segment id prediction function
-int vp9_get_pred_mi_segid(VP9_COMMON *cm, BLOCK_SIZE_TYPE sb_type,
-                          uint8_t *segment_ids, int mi_row, int mi_col) {
-  const int mi_index = mi_row * cm->mi_cols + mi_col;
-  const int bw = 1 << mi_width_log2(sb_type);
-  const int bh = 1 << mi_height_log2(sb_type);
-  const int ymis = MIN(cm->mi_rows - mi_row, bh);
+int vp9_get_segment_id(VP9_COMMON *cm, const uint8_t *segment_ids,
+                       BLOCK_SIZE_TYPE bsize, int mi_row, int mi_col) {
+  const int mi_offset = mi_row * cm->mi_cols + mi_col;
+  const int bw = 1 << mi_width_log2(bsize);
+  const int bh = 1 << mi_height_log2(bsize);
   const int xmis = MIN(cm->mi_cols - mi_col, bw);
-  int segment_id = INT_MAX;
-  int x, y;
+  const int ymis = MIN(cm->mi_rows - mi_row, bh);
+  int x, y, segment_id = INT_MAX;
 
-  for (y = 0; y < ymis; y++) {
-    for (x = 0; x < xmis; x++) {
-      const int index = mi_index + (y * cm->mi_cols + x);
-      segment_id = MIN(segment_id, segment_ids[index]);
-    }
-  }
+  for (y = 0; y < ymis; y++)
+    for (x = 0; x < xmis; x++)
+      segment_id = MIN(segment_id,
+                       segment_ids[mi_offset + y * cm->mi_cols + x]);
+
+  assert(segment_id >= 0 && segment_id < MAX_MB_SEGMENTS);
   return segment_id;
 }
diff --git a/vp9/common/vp9_pred_common.h b/vp9/common/vp9_pred_common.h
index 100102b61..2717ff188 100644
--- a/vp9/common/vp9_pred_common.h
+++ b/vp9/common/vp9_pred_common.h
@@ -27,27 +27,21 @@ typedef enum {
   PRED_TX_SIZE = 8
 } PRED_ID;
 
-unsigned char vp9_get_pred_context(const VP9_COMMON *const cm,
-                                   const MACROBLOCKD *const xd,
+unsigned char vp9_get_pred_context(const VP9_COMMON *cm, const MACROBLOCKD *xd,
                                    PRED_ID pred_id);
 
-vp9_prob vp9_get_pred_prob(const VP9_COMMON *const cm,
-                           const MACROBLOCKD *const xd,
+vp9_prob vp9_get_pred_prob(const VP9_COMMON *cm, const MACROBLOCKD *xd,
                            PRED_ID pred_id);
 
-const vp9_prob *vp9_get_pred_probs(const VP9_COMMON *const cm,
-                                   const MACROBLOCKD *const xd,
+const vp9_prob *vp9_get_pred_probs(const VP9_COMMON *cm, const MACROBLOCKD *xd,
                                    PRED_ID pred_id);
 
-unsigned char vp9_get_pred_flag(const MACROBLOCKD *const xd,
-                                PRED_ID pred_id);
+unsigned char vp9_get_pred_flag(const MACROBLOCKD *xd, PRED_ID pred_id);
 
-void vp9_set_pred_flag(MACROBLOCKD *const xd,
-                       PRED_ID pred_id,
+void vp9_set_pred_flag(MACROBLOCKD *xd, PRED_ID pred_id,
                        unsigned char pred_flag);
 
-
-int vp9_get_pred_mi_segid(VP9_COMMON *cm, BLOCK_SIZE_TYPE sb_type,
-                          uint8_t *segment_ids, int mi_row, int mi_col);
+int vp9_get_segment_id(VP9_COMMON *cm, const uint8_t *segment_ids,
+                       BLOCK_SIZE_TYPE bsize, int mi_row, int mi_col);
 
 #endif  // VP9_COMMON_VP9_PRED_COMMON_H_
diff --git a/vp9/decoder/vp9_decodemv.c b/vp9/decoder/vp9_decodemv.c
index 7ca69d0a5..020e9c6cc 100644
--- a/vp9/decoder/vp9_decodemv.c
+++ b/vp9/decoder/vp9_decodemv.c
@@ -43,7 +43,7 @@ static MB_PREDICTION_MODE read_intra_mode(vp9_reader *r, const vp9_prob *p) {
   return treed_read(r, vp9_intra_mode_tree, p);
 }
 
-static int read_mb_segid(vp9_reader *r, MACROBLOCKD *xd) {
+static int read_segment_id(vp9_reader *r, MACROBLOCKD *xd) {
   return treed_read(r, vp9_segment_tree, xd->mb_segment_tree_probs);
 }
 
@@ -86,21 +86,34 @@ static TX_SIZE get_txfm_size(VP9D_COMP *pbi, TXFM_MODE txfm_mode,
     return TX_4X4;
 }
 
-static void set_segment_id(VP9_COMMON *cm, MB_MODE_INFO *mbmi,
+static void set_segment_id(VP9_COMMON *cm, BLOCK_SIZE_TYPE bsize,
                            int mi_row, int mi_col, int segment_id) {
-  const int mi_index = mi_row * cm->mi_cols + mi_col;
-  const BLOCK_SIZE_TYPE sb_type = mbmi->sb_type;
-  const int bw = 1 << mi_width_log2(sb_type);
-  const int bh = 1 << mi_height_log2(sb_type);
-  const int ymis = MIN(cm->mi_rows - mi_row, bh);
+  const int mi_offset = mi_row * cm->mi_cols + mi_col;
+  const int bw = 1 << mi_width_log2(bsize);
+  const int bh = 1 << mi_height_log2(bsize);
   const int xmis = MIN(cm->mi_cols - mi_col, bw);
+  const int ymis = MIN(cm->mi_rows - mi_row, bh);
   int x, y;
 
-  for (y = 0; y < ymis; y++) {
-    for (x = 0; x < xmis; x++) {
-      const int index = mi_index + (y * cm->mi_cols + x);
-      cm->last_frame_seg_map[index] = segment_id;
-    }
+  assert(segment_id >= 0 && segment_id < MAX_MB_SEGMENTS);
+
+  for (y = 0; y < ymis; y++)
+    for (x = 0; x < xmis; x++)
+      cm->last_frame_seg_map[mi_offset + y * cm->mi_cols + x] = segment_id;
+}
+
+static int read_intra_segment_id(VP9D_COMP *pbi, int mi_row, int mi_col,
+                                 vp9_reader *r) {
+  VP9_COMMON *const cm = &pbi->common;
+  MACROBLOCKD *const xd = &pbi->mb;
+  const BLOCK_SIZE_TYPE bsize = xd->mode_info_context->mbmi.sb_type;
+
+  if (xd->segmentation_enabled && xd->update_mb_segmentation_map) {
+    const int segment_id = read_segment_id(r, xd);
+    set_segment_id(cm, bsize, mi_row, mi_col, segment_id);
+    return segment_id;
+  } else {
+    return 0;
   }
 }
 
@@ -111,13 +124,7 @@ static void kfread_modes(VP9D_COMP *pbi, MODE_INFO *m,
   MACROBLOCKD *const xd = &pbi->mb;
   const int mis = cm->mode_info_stride;
 
-  // Read segmentation map if it is being updated explicitly this frame
-  m->mbmi.segment_id = 0;
-  if (xd->segmentation_enabled && xd->update_mb_segmentation_map) {
-    m->mbmi.segment_id = read_mb_segid(r, xd);
-    set_segment_id(cm, &m->mbmi, mi_row, mi_col, m->mbmi.segment_id);
-  }
-
+  m->mbmi.segment_id = read_intra_segment_id(pbi, mi_row, mi_col, r);
   m->mbmi.mb_skip_coeff = vp9_segfeature_active(xd, m->mbmi.segment_id,
                                                 SEG_LVL_SKIP);
   if (!m->mbmi.mb_skip_coeff) {
@@ -388,46 +395,32 @@ static void mb_mode_mv_init(VP9D_COMP *pbi, vp9_reader *r) {
   }
 }
 
-// This function either reads the segment id for the current macroblock from
-// the bitstream or if the value is temporally predicted asserts the predicted
-// value
-static int read_mb_segment_id(VP9D_COMP *pbi, int mi_row, int mi_col,
-                              vp9_reader *r) {
+static int read_inter_segment_id(VP9D_COMP *pbi, int mi_row, int mi_col,
+                                 vp9_reader *r) {
   VP9_COMMON *const cm = &pbi->common;
   MACROBLOCKD *const xd = &pbi->mb;
-  MODE_INFO *const mi = xd->mode_info_context;
-  MB_MODE_INFO *const mbmi = &mi->mbmi;
+  const BLOCK_SIZE_TYPE bsize = xd->mode_info_context->mbmi.sb_type;
+  const int pred_segment_id = vp9_get_segment_id(cm, cm->last_frame_seg_map,
+                                                 bsize, mi_row, mi_col);
+  int segment_id;
 
   if (!xd->segmentation_enabled)
     return 0;  // Default for disabled segmentation
 
-  if (xd->update_mb_segmentation_map) {
-    int segment_id;
-
-    if (cm->temporal_update) {
-      // Temporal coding of the segment id for this mb is enabled.
-      // Get the context based probability for reading the
-      // prediction status flag
-      const vp9_prob pred_prob = vp9_get_pred_prob(cm, xd, PRED_SEG_ID);
-      const int pred_flag = vp9_read(r, pred_prob);
-      vp9_set_pred_flag(xd, PRED_SEG_ID, pred_flag);
-
-      // If the value is flagged as correctly predicted
-      // then use the predicted value, otherwise decode it explicitly
-      segment_id = pred_flag ? vp9_get_pred_mi_segid(cm, mbmi->sb_type,
-                                                     cm->last_frame_seg_map,
-                                                     mi_row, mi_col)
-                             : read_mb_segid(r, xd);
-    } else {
-      segment_id = read_mb_segid(r, xd);  // Normal unpredicted coding mode
-    }
+  if (!xd->update_mb_segmentation_map)
+    return pred_segment_id;
 
-    set_segment_id(cm, mbmi, mi_row, mi_col, segment_id);  // Side effect
-    return segment_id;
+  if (cm->temporal_update) {
+    const vp9_prob pred_prob = vp9_get_pred_prob(cm, xd, PRED_SEG_ID);
+    const int pred_flag = vp9_read(r, pred_prob);
+    vp9_set_pred_flag(xd, PRED_SEG_ID, pred_flag);
+    segment_id = pred_flag ? pred_segment_id
+                           : read_segment_id(r, xd);
   } else {
-    return vp9_get_pred_mi_segid(cm, mbmi->sb_type, cm->last_frame_seg_map,
-                                 mi_row, mi_col);
+    segment_id = read_segment_id(r, xd);
   }
+  set_segment_id(cm, bsize, mi_row, mi_col, segment_id);
+  return segment_id;
 }
 
 
@@ -559,7 +552,7 @@ static void read_mb_modes_mv(VP9D_COMP *pbi, MODE_INFO *mi, MB_MODE_INFO *mbmi,
   mb_to_right_edge = xd->mb_to_right_edge + RIGHT_BOTTOM_MARGIN;
 
   // Read the macroblock segment id.
-  mbmi->segment_id = read_mb_segment_id(pbi, mi_row, mi_col, r);
+  mbmi->segment_id = read_inter_segment_id(pbi, mi_row, mi_col, r);
 
   mbmi->mb_skip_coeff = vp9_segfeature_active(xd, mbmi->segment_id,
                                               SEG_LVL_SKIP);
diff --git a/vp9/encoder/vp9_bitstream.c b/vp9/encoder/vp9_bitstream.c
index 7538151c6..decd5a2f3 100644
--- a/vp9/encoder/vp9_bitstream.c
+++ b/vp9/encoder/vp9_bitstream.c
@@ -562,13 +562,11 @@ static void write_sb_mv_ref(vp9_writer *bc, MB_PREDICTION_MODE m,
               vp9_sb_mv_ref_encoding_array - NEARESTMV + m);
 }
 
-// This function writes the current macro block's segnment id to the bitstream
-// It should only be called if a segment map update is indicated.
-static void write_mb_segid(vp9_writer *bc,
-                           const MB_MODE_INFO *mi, const MACROBLOCKD *xd) {
+
+static void write_segment_id(vp9_writer *w, const MACROBLOCKD *xd,
+                             int segment_id) {
   if (xd->segmentation_enabled && xd->update_mb_segmentation_map)
-    treed_write(bc, vp9_segment_tree, xd->mb_segment_tree_probs,
-                mi->segment_id, 3);
+    treed_write(w, vp9_segment_tree, xd->mb_segment_tree_probs, segment_id, 3);
 }
 
 // This function encodes the reference frame
@@ -643,10 +641,10 @@ static void pack_inter_mode_mvs(VP9_COMP *cpi, MODE_INFO *m,
 
       // If the mb segment id wasn't predicted code explicitly
       if (!prediction_flag)
-        write_mb_segid(bc, mi, &cpi->mb.e_mbd);
+        write_segment_id(bc, xd, mi->segment_id);
     } else {
       // Normal unpredicted coding
-      write_mb_segid(bc, mi, &cpi->mb.e_mbd);
+        write_segment_id(bc, xd, mi->segment_id);
     }
   }
 
@@ -782,7 +780,7 @@ static void write_mb_modes_kf(const VP9_COMP *cpi,
   int skip_coeff;
 
   if (xd->update_mb_segmentation_map)
-    write_mb_segid(bc, &m->mbmi, xd);
+    write_segment_id(bc, xd, m->mbmi.segment_id);
 
   if (vp9_segfeature_active(xd, segment_id, SEG_LVL_SKIP)) {
     skip_coeff = 1;
diff --git a/vp9/encoder/vp9_encodeframe.c b/vp9/encoder/vp9_encodeframe.c
index 4b1ff103a..168b46012 100644
--- a/vp9/encoder/vp9_encodeframe.c
+++ b/vp9/encoder/vp9_encodeframe.c
@@ -516,9 +516,8 @@ static void set_offsets(VP9_COMP *cpi, int mi_row, int mi_col,
   if (xd->segmentation_enabled) {
     uint8_t *map = xd->update_mb_segmentation_map ? cpi->segmentation_map
                                                   : cm->last_frame_seg_map;
-    mbmi->segment_id = vp9_get_pred_mi_segid(cm, bsize, map, mi_row, mi_col);
+    mbmi->segment_id = vp9_get_segment_id(cm, map, bsize, mi_row, mi_col);
 
-    assert(mbmi->segment_id <= (MAX_MB_SEGMENTS-1));
     vp9_mb_init_quantizer(cpi, x);
 
     if (xd->segmentation_enabled && cpi->seg0_cnt > 0
diff --git a/vp9/encoder/vp9_segmentation.c b/vp9/encoder/vp9_segmentation.c
index af218b7d6..eaa1d534a 100644
--- a/vp9/encoder/vp9_segmentation.c
+++ b/vp9/encoder/vp9_segmentation.c
@@ -115,8 +115,7 @@ static int cost_segmap(MACROBLOCKD *xd, int *segcounts, vp9_prob *probs) {
   return cost;
 }
 
-static void count_segs(VP9_COMP *cpi,
-                       MODE_INFO *mi,
+static void count_segs(VP9_COMP *cpi, MODE_INFO *mi,
                        int *no_pred_segcounts,
                        int (*temporal_predictor_count)[2],
                        int *t_unpred_seg_counts,
@@ -138,20 +137,18 @@ static void count_segs(VP9_COMP *cpi,
   // Temporal prediction not allowed on key frames
   if (cm->frame_type != KEY_FRAME) {
     // Test to see if the segment id matches the predicted value.
-    const int pred_seg_id = vp9_get_pred_mi_segid(cm, mi->mbmi.sb_type,
-                                                  cm->last_frame_seg_map,
-                                                  mi_row, mi_col);
-    const int seg_predicted = (segment_id == pred_seg_id);
-
-    // Get the segment id prediction context
+    const int pred_segment_id = vp9_get_segment_id(cm, cm->last_frame_seg_map,
+                                                   mi->mbmi.sb_type,
+                                                   mi_row, mi_col);
+    const int pred_flag = pred_segment_id == segment_id;
     const int pred_context = vp9_get_pred_context(cm, xd, PRED_SEG_ID);
 
     // Store the prediction status for this mb and update counts
     // as appropriate
-    vp9_set_pred_flag(xd, PRED_SEG_ID, seg_predicted);
-    temporal_predictor_count[pred_context][seg_predicted]++;
+    vp9_set_pred_flag(xd, PRED_SEG_ID, pred_flag);
+    temporal_predictor_count[pred_context][pred_flag]++;
 
-    if (!seg_predicted)
+    if (!pred_flag)
       // Update the "unpredicted" segment count
       t_unpred_seg_counts[segment_id]++;
   }