Cleaning up entropy probability update in encoder.

Change-Id: I94cb9e3d910dff74bf90906dd96e3a4e06ebdbe6

2 files changed, 34 insertions, 70 deletions

diff --git a/vp9/encoder/vp9_bitstream.c b/vp9/encoder/vp9_bitstream.c
index efbadbab0..9bd0123a8 100644
--- a/vp9/encoder/vp9_bitstream.c
+++ b/vp9/encoder/vp9_bitstream.c
@@ -163,27 +163,17 @@ void vp9_encode_unsigned_max(struct vp9_write_bit_buffer *wb,
   vp9_wb_write_literal(wb, data, get_unsigned_bits(max));
 }
 
-static void update_mode(vp9_writer *w, int n, vp9_tree tree,
-                        vp9_prob Pcur[/* n-1 */],
-                        unsigned int bct[/* n-1 */][2],
-                        const unsigned int num_events[/* n */]) {
-  int i = 0;
+static void prob_diff_update(const vp9_tree_index *tree,
+                             vp9_prob probs[/*n - 1*/],
+                             const unsigned int counts[/*n - 1*/],
+                             int n, vp9_writer *w) {
+  int i;
+  unsigned int branch_ct[32][2];
+  assert(n <= 32);
 
-  vp9_tree_probs_from_distribution(tree, bct, num_events);
+  vp9_tree_probs_from_distribution(tree, branch_ct, counts);
   for (i = 0; i < n - 1; ++i)
-    vp9_cond_prob_diff_update(w, &Pcur[i], bct[i]);
-}
-
-static void update_mbintra_mode_probs(VP9_COMP* const cpi,
-                                      vp9_writer* const bc) {
-  VP9_COMMON *const cm = &cpi->common;
-  int j;
-  unsigned int bct[INTRA_MODES - 1][2];
-
-  for (j = 0; j < BLOCK_SIZE_GROUPS; j++)
-    update_mode(bc, INTRA_MODES, vp9_intra_mode_tree,
-                cm->fc.y_mode_prob[j], bct,
-                (unsigned int *)cpi->y_mode_count[j]);
+    vp9_cond_prob_diff_update(w, &probs[i], branch_ct[i]);
 }
 
 static void write_selected_tx_size(const VP9_COMP *cpi, MODE_INFO *m,
@@ -227,16 +217,11 @@ static void write_intra_mode(vp9_writer *bc, int m, const vp9_prob *p) {
 
 static void update_switchable_interp_probs(VP9_COMP *cpi, vp9_writer *w) {
   VP9_COMMON *const cm = &cpi->common;
-  unsigned int branch_ct[SWITCHABLE_FILTERS - 1][2];
-  int i, j;
-  for (j = 0; j < SWITCHABLE_FILTER_CONTEXTS; ++j) {
-    vp9_tree_probs_from_distribution(vp9_switchable_interp_tree, branch_ct,
-                                     cm->counts.switchable_interp[j]);
-
-    for (i = 0; i < SWITCHABLE_FILTERS - 1; ++i)
-      vp9_cond_prob_diff_update(w, &cm->fc.switchable_interp_prob[j][i],
-                                branch_ct[i]);
-  }
+  int j;
+  for (j = 0; j < SWITCHABLE_FILTER_CONTEXTS; ++j)
+    prob_diff_update(vp9_switchable_interp_tree,
+                     cm->fc.switchable_interp_prob[j],
+                     cm->counts.switchable_interp[j], SWITCHABLE_FILTERS, w);
 
 #ifdef MODE_STATS
   if (!cpi->dummy_packing)
@@ -244,20 +229,6 @@ static void update_switchable_interp_probs(VP9_COMP *cpi, vp9_writer *w) {
 #endif
 }
 
-static void update_inter_mode_probs(VP9_COMMON *cm, vp9_writer *w) {
-  int i, j;
-
-  for (i = 0; i < INTER_MODE_CONTEXTS; ++i) {
-    unsigned int branch_ct[INTER_MODES - 1][2];
-    vp9_tree_probs_from_distribution(vp9_inter_mode_tree, branch_ct,
-                                     cm->counts.inter_mode[i]);
-
-    for (j = 0; j < INTER_MODES - 1; ++j)
-      vp9_cond_prob_diff_update(w, &cm->fc.inter_mode_probs[i][j],
-                                branch_ct[j]);
-  }
-}
-
 static void pack_mb_tokens(vp9_writer* const w,
                            TOKENEXTRA **tp,
                            const TOKENEXTRA *const stop) {
@@ -1382,7 +1353,10 @@ static size_t write_compressed_header(VP9_COMP *cpi, uint8_t *data) {
     active_section = 1;
 #endif
 
-    update_inter_mode_probs(cm, &header_bc);
+    for (i = 0; i < INTER_MODE_CONTEXTS; ++i)
+      prob_diff_update(vp9_inter_mode_tree, cm->fc.inter_mode_probs[i],
+                       cm->counts.inter_mode[i], INTER_MODES, &header_bc);
+
     vp9_zero(cm->counts.inter_mode);
 
     if (cm->mcomp_filter_type == SWITCHABLE)
@@ -1421,14 +1395,15 @@ static size_t write_compressed_header(VP9_COMP *cpi, uint8_t *data) {
         vp9_cond_prob_diff_update(&header_bc, &fc->comp_ref_prob[i],
                                   cpi->comp_ref_count[i]);
 
-    update_mbintra_mode_probs(cpi, &header_bc);
+    for (i = 0; i < BLOCK_SIZE_GROUPS; ++i)
+      prob_diff_update(vp9_intra_mode_tree, cm->fc.y_mode_prob[i],
+                       (unsigned int *)cpi->y_mode_count[i], INTRA_MODES,
+                       &header_bc);
 
-    for (i = 0; i < PARTITION_CONTEXTS; ++i) {
-      unsigned int bct[PARTITION_TYPES - 1][2];
-      update_mode(&header_bc, PARTITION_TYPES, vp9_partition_tree,
-                  fc->partition_prob[i], bct,
-                  (unsigned int *)cpi->partition_count[i]);
-    }
+    for (i = 0; i < PARTITION_CONTEXTS; ++i)
+      prob_diff_update(vp9_partition_tree, fc->partition_prob[i],
+                       (unsigned int *)cpi->partition_count[i], PARTITION_TYPES,
+                       &header_bc);
 
     vp9_write_nmv_probs(cpi, cm->allow_high_precision_mv, &header_bc);
   }
diff --git a/vp9/encoder/vp9_encodemv.c b/vp9/encoder/vp9_encodemv.c
index 7e838c901..bb9baff44 100644
--- a/vp9/encoder/vp9_encodemv.c
+++ b/vp9/encoder/vp9_encodemv.c
@@ -153,11 +153,8 @@ static void counts_to_nmv_context(
   vp9_tree_probs_from_distribution(vp9_mv_joint_tree, branch_ct_joint,
                                    nmv_count->joints);
   for (i = 0; i < 2; ++i) {
-    const uint32_t s0 = nmv_count->comps[i].sign[0];
-    const uint32_t s1 = nmv_count->comps[i].sign[1];
-
-    branch_ct_sign[i][0] = s0;
-    branch_ct_sign[i][1] = s1;
+    branch_ct_sign[i][0] = nmv_count->comps[i].sign[0];
+    branch_ct_sign[i][1] = nmv_count->comps[i].sign[1];
     vp9_tree_probs_from_distribution(vp9_mv_class_tree,
                                     branch_ct_classes[i],
                                     nmv_count->comps[i].classes);
@@ -165,11 +162,8 @@ static void counts_to_nmv_context(
                                      branch_ct_class0[i],
                                      nmv_count->comps[i].class0);
     for (j = 0; j < MV_OFFSET_BITS; ++j) {
-      const uint32_t b0 = nmv_count->comps[i].bits[j][0];
-      const uint32_t b1 = nmv_count->comps[i].bits[j][1];
-
-      branch_ct_bits[i][j][0] = b0;
-      branch_ct_bits[i][j][1] = b1;
+      branch_ct_bits[i][j][0] = nmv_count->comps[i].bits[j][0];
+      branch_ct_bits[i][j][1] = nmv_count->comps[i].bits[j][1];
     }
   }
   for (i = 0; i < 2; ++i) {
@@ -184,16 +178,11 @@ static void counts_to_nmv_context(
   }
   if (usehp) {
     for (i = 0; i < 2; ++i) {
-      const uint32_t c0_hp0 = nmv_count->comps[i].class0_hp[0];
-      const uint32_t c0_hp1 = nmv_count->comps[i].class0_hp[1];
-      const uint32_t hp0 = nmv_count->comps[i].hp[0];
-      const uint32_t hp1 = nmv_count->comps[i].hp[1];
-
-      branch_ct_class0_hp[i][0] = c0_hp0;
-      branch_ct_class0_hp[i][1] = c0_hp1;
+      branch_ct_class0_hp[i][0] = nmv_count->comps[i].class0_hp[0];
+      branch_ct_class0_hp[i][1] = nmv_count->comps[i].class0_hp[1];
 
-      branch_ct_hp[i][0] = hp0;
-      branch_ct_hp[i][1] = hp1;
+      branch_ct_hp[i][0] = nmv_count->comps[i].hp[0];
+      branch_ct_hp[i][1] = nmv_count->comps[i].hp[1];
     }
   }
 }