Add av1_get_br_level_counts()

which calculates br level counts in block.

Change-Id: Iec5c9a8a0ce7e57bf5fb29db552f6a9065c5de88
diff --git a/av1/common/txb_common.c b/av1/common/txb_common.c
index a5c8b7d..a4dc779 100644
--- a/av1/common/txb_common.c
+++ b/av1/common/txb_common.c
@@ -296,3 +296,17 @@
     }
   }
 }
+
+void av1_get_br_level_counts(const uint8_t *const levels, const int width,
+                             const int height, uint8_t *const level_counts) {
+  const int stride = width + TX_PAD_HOR;
+  const int level_minus_1 = NUM_BASE_LEVELS;
+
+  for (int row = 0; row < height; ++row) {
+    for (int col = 0; col < width; ++col) {
+      level_counts[row * width + col] =
+          get_level_count(levels, stride, row, col, level_minus_1,
+                          br_ref_offset, BR_CONTEXT_POSITION_NUM);
+    }
+  }
+}
diff --git a/av1/common/txb_common.h b/av1/common/txb_common.h
index 590228b..21b4abc 100644
--- a/av1/common/txb_common.h
+++ b/av1/common/txb_common.h
@@ -281,15 +281,12 @@
 
 static INLINE int get_br_ctx(const uint8_t *const levels,
                              const int c,  // raster order
-                             const int bwl) {
+                             const int bwl, const int count) {
   const int row = c >> bwl;
   const int col = c - (row << bwl);
   const int stride = (1 << bwl) + TX_PAD_HOR;
-  const int level_minus_1 = NUM_BASE_LEVELS;
   int mag = 0;
   int nb_mag[3] = { 0 };
-  const int count = get_level_count(levels, stride, row, col, level_minus_1,
-                                    br_ref_offset, BR_CONTEXT_POSITION_NUM);
   get_level_mag(levels, stride, row, col, nb_mag);
   for (int idx = 0; idx < 3; ++idx) mag = AOMMAX(mag, nb_mag[idx]);
   const int ctx = get_br_ctx_from_count_mag(row, col, count, mag);
@@ -571,4 +568,7 @@
                                const int level_minus_1, const int width,
                                const int height, uint8_t *const level_counts);
 
+void av1_get_br_level_counts(const uint8_t *const levels, const int width,
+                             const int height, uint8_t *const level_counts);
+
 #endif  // AV1_COMMON_TXB_COMMON_H_
diff --git a/av1/decoder/decodetxb.c b/av1/decoder/decodetxb.c
index 1c9f3b1..e0a1c76 100644
--- a/av1/decoder/decodetxb.c
+++ b/av1/decoder/decodetxb.c
@@ -74,6 +74,7 @@
   int cul_level = 0;
   uint8_t levels_buf[TX_PAD_2D];
   uint8_t *const levels = set_levels(levels_buf, width);
+  uint8_t level_counts[MAX_TX_SQUARE];
   int8_t signs[MAX_TX_SQUARE];
 
   memset(tcoeffs, 0, sizeof(*tcoeffs) * seg_eob);
@@ -205,7 +206,6 @@
 #if USE_CAUSAL_BASE_CTX
   update_eob = *eob - 1;
 #else
-  uint8_t level_counts[MAX_TX_SQUARE];
   int i;
   for (i = 0; i < NUM_BASE_LEVELS; ++i) {
     av1_get_base_level_counts(levels, i, width, height, level_counts);
@@ -255,52 +255,55 @@
     }
   }
 
-  for (c = update_eob; c >= 0; --c) {
-    uint8_t *const level = &levels[get_paded_idx(scan[c], bwl)];
-    int idx;
-    int ctx;
+  if (update_eob >= 0) {
+    av1_get_br_level_counts(levels, width, height, level_counts);
+    for (c = update_eob; c >= 0; --c) {
+      uint8_t *const level = &levels[get_paded_idx(scan[c], bwl)];
+      int idx;
+      int ctx;
 
-    if (*level <= NUM_BASE_LEVELS) continue;
+      if (*level <= NUM_BASE_LEVELS) continue;
 
-    ctx = get_br_ctx(levels, scan[c], bwl);
+      ctx = get_br_ctx(levels, scan[c], bwl, level_counts[scan[c]]);
 
-    for (idx = 0; idx < BASE_RANGE_SETS; ++idx) {
-      if (av1_read_record_bin(
-              counts, r, ec_ctx->coeff_br_cdf[txs_ctx][plane_type][idx][ctx], 2,
-              ACCT_STR)) {
-        int extra_bits = (1 << br_extra_bits[idx]) - 1;
-        //        int br_offset = aom_read_literal(r, extra_bits, ACCT_STR);
-        int br_offset = 0;
-        int tok;
-        if (counts) ++counts->coeff_br[txs_ctx][plane_type][idx][ctx][1];
-        for (tok = 0; tok < extra_bits; ++tok) {
-          if (av1_read_record_bin(
-                  counts, r, ec_ctx->coeff_lps_cdf[txs_ctx][plane_type][ctx], 2,
-                  ACCT_STR)) {
-            br_offset = tok;
-            if (counts) ++counts->coeff_lps[txs_ctx][plane_type][ctx][1];
-            break;
+      for (idx = 0; idx < BASE_RANGE_SETS; ++idx) {
+        if (av1_read_record_bin(
+                counts, r, ec_ctx->coeff_br_cdf[txs_ctx][plane_type][idx][ctx],
+                2, ACCT_STR)) {
+          int extra_bits = (1 << br_extra_bits[idx]) - 1;
+          //        int br_offset = aom_read_literal(r, extra_bits, ACCT_STR);
+          int br_offset = 0;
+          int tok;
+          if (counts) ++counts->coeff_br[txs_ctx][plane_type][idx][ctx][1];
+          for (tok = 0; tok < extra_bits; ++tok) {
+            if (av1_read_record_bin(
+                    counts, r, ec_ctx->coeff_lps_cdf[txs_ctx][plane_type][ctx],
+                    2, ACCT_STR)) {
+              br_offset = tok;
+              if (counts) ++counts->coeff_lps[txs_ctx][plane_type][ctx][1];
+              break;
+            }
+            if (counts) ++counts->coeff_lps[txs_ctx][plane_type][ctx][0];
           }
-          if (counts) ++counts->coeff_lps[txs_ctx][plane_type][ctx][0];
+          if (tok == extra_bits) br_offset = extra_bits;
+
+          int br_base = br_index_to_coeff[idx];
+
+          *level = NUM_BASE_LEVELS + 1 + br_base + br_offset;
+          cul_level += *level;
+          break;
         }
-        if (tok == extra_bits) br_offset = extra_bits;
-
-        int br_base = br_index_to_coeff[idx];
-
-        *level = NUM_BASE_LEVELS + 1 + br_base + br_offset;
-        cul_level += *level;
-        break;
+        if (counts) ++counts->coeff_br[txs_ctx][plane_type][idx][ctx][0];
       }
-      if (counts) ++counts->coeff_br[txs_ctx][plane_type][idx][ctx][0];
+
+      if (idx < BASE_RANGE_SETS) continue;
+
+      // decode 0-th order Golomb code
+      *level = COEFF_BASE_RANGE + 1 + NUM_BASE_LEVELS;
+      // Save golomb in tcoeffs because adding it to level may incur overflow
+      tcoeffs[scan[c]] = read_golomb(xd, r, counts);
+      cul_level += *level + tcoeffs[scan[c]];
     }
-
-    if (idx < BASE_RANGE_SETS) continue;
-
-    // decode 0-th order Golomb code
-    *level = COEFF_BASE_RANGE + 1 + NUM_BASE_LEVELS;
-    // Save golomb in tcoeffs because adding it to level may incur overflow
-    tcoeffs[scan[c]] = read_golomb(xd, r, counts);
-    cul_level += *level + tcoeffs[scan[c]];
   }
 
   for (c = 0; c < *eob; ++c) {
diff --git a/av1/encoder/encodetxb.c b/av1/encoder/encodetxb.c
index 5fd2836..2ac2eac 100644
--- a/av1/encoder/encodetxb.c
+++ b/av1/encoder/encodetxb.c
@@ -193,9 +193,9 @@
                           const TxbInfo *const txb_info,
                           const LV_MAP_COEFF_COST *const txb_costs);
 
-void get_dist_cost_stats(LevelDownStats *stats, int scan_idx,
-                         const LV_MAP_COEFF_COST *txb_costs,
-                         TxbInfo *txb_info) {
+static void get_dist_cost_stats(LevelDownStats *stats, int scan_idx,
+                                const LV_MAP_COEFF_COST *txb_costs,
+                                TxbInfo *txb_info) {
   const int16_t *scan = txb_info->scan_order->scan;
   const int coeff_idx = scan[scan_idx];
   const tran_low_t qc = txb_info->qcoeff[coeff_idx];
@@ -297,6 +297,7 @@
   FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
   uint8_t levels_buf[TX_PAD_2D];
   uint8_t *const levels = set_levels(levels_buf, width);
+  uint8_t level_counts[MAX_TX_SQUARE];
 
   (void)blk_row;
   (void)blk_col;
@@ -371,7 +372,6 @@
   }
   update_eob = eob - 1;
 #else
-  uint8_t level_counts[MAX_TX_SQUARE];
   for (int i = 1; i < eob; ++i) {
     c = eob - 1 - i;
     int coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type);
@@ -426,52 +426,55 @@
     }
   }
 
-  for (c = update_eob; c >= 0; --c) {
-    const tran_low_t level = abs(tcoeff[scan[c]]);
-    int idx;
-    int ctx;
+  if (update_eob >= 0) {
+    av1_get_br_level_counts(levels, width, height, level_counts);
+    for (c = update_eob; c >= 0; --c) {
+      const tran_low_t level = abs(tcoeff[scan[c]]);
+      int idx;
+      int ctx;
 
-    if (level <= NUM_BASE_LEVELS) continue;
+      if (level <= NUM_BASE_LEVELS) continue;
 
-    // level is above 1.
-    ctx = get_br_ctx(levels, scan[c], bwl);
+      // level is above 1.
+      ctx = get_br_ctx(levels, scan[c], bwl, level_counts[scan[c]]);
 
-    int base_range = level - 1 - NUM_BASE_LEVELS;
-    int br_set_idx = 0;
-    int br_base = 0;
-    int br_offset = 0;
+      int base_range = level - 1 - NUM_BASE_LEVELS;
+      int br_set_idx = 0;
+      int br_base = 0;
+      int br_offset = 0;
 
-    if (base_range >= COEFF_BASE_RANGE)
-      br_set_idx = BASE_RANGE_SETS;
-    else
-      br_set_idx = coeff_to_br_index[base_range];
+      if (base_range >= COEFF_BASE_RANGE)
+        br_set_idx = BASE_RANGE_SETS;
+      else
+        br_set_idx = coeff_to_br_index[base_range];
 
-    for (idx = 0; idx < BASE_RANGE_SETS; ++idx) {
-      aom_write_bin(w, idx == br_set_idx,
-                    ec_ctx->coeff_br_cdf[txs_ctx][plane_type][idx][ctx], 2);
-      if (idx == br_set_idx) {
-        br_base = br_index_to_coeff[br_set_idx];
-        br_offset = base_range - br_base;
-        int extra_bits = (1 << br_extra_bits[idx]) - 1;
-        for (int tok = 0; tok < extra_bits; ++tok) {
-          if (tok == br_offset) {
-            aom_write_bin(w, 1, ec_ctx->coeff_lps_cdf[txs_ctx][plane_type][ctx],
+      for (idx = 0; idx < BASE_RANGE_SETS; ++idx) {
+        aom_write_bin(w, idx == br_set_idx,
+                      ec_ctx->coeff_br_cdf[txs_ctx][plane_type][idx][ctx], 2);
+        if (idx == br_set_idx) {
+          br_base = br_index_to_coeff[br_set_idx];
+          br_offset = base_range - br_base;
+          int extra_bits = (1 << br_extra_bits[idx]) - 1;
+          for (int tok = 0; tok < extra_bits; ++tok) {
+            if (tok == br_offset) {
+              aom_write_bin(w, 1,
+                            ec_ctx->coeff_lps_cdf[txs_ctx][plane_type][ctx], 2);
+              break;
+            }
+            aom_write_bin(w, 0, ec_ctx->coeff_lps_cdf[txs_ctx][plane_type][ctx],
                           2);
-            break;
           }
-          aom_write_bin(w, 0, ec_ctx->coeff_lps_cdf[txs_ctx][plane_type][ctx],
-                        2);
+          //        aom_write_literal(w, br_offset, br_extra_bits[idx]);
+          break;
         }
-        //        aom_write_literal(w, br_offset, br_extra_bits[idx]);
-        break;
       }
+
+      if (br_set_idx < BASE_RANGE_SETS) continue;
+
+      // use 0-th order Golomb code to handle the residual level.
+      write_golomb(
+          w, abs(tcoeff[scan[c]]) - COEFF_BASE_RANGE - 1 - NUM_BASE_LEVELS);
     }
-
-    if (br_set_idx < BASE_RANGE_SETS) continue;
-
-    // use 0-th order Golomb code to handle the residual level.
-    write_golomb(w,
-                 abs(tcoeff[scan[c]]) - COEFF_BASE_RANGE - 1 - NUM_BASE_LEVELS);
   }
 }
 
@@ -590,6 +593,7 @@
   const int16_t *scan = scan_order->scan;
   uint8_t levels_buf[TX_PAD_2D];
   uint8_t *const levels = set_levels(levels_buf, width);
+  uint8_t level_counts[MAX_TX_SQUARE];
 
   LV_MAP_COEFF_COST *coeff_costs = &x->coeff_costs[txs_ctx][plane_type];
 
@@ -610,6 +614,7 @@
   const int seg_eob = tx_size_2d[tx_size];
   int eob_cost = get_eob_cost(eob, seg_eob, coeff_costs, tx_type);
 
+  av1_get_br_level_counts(levels, width, height, level_counts);
   cost += eob_cost;
 #if USE_CAUSAL_BASE_CTX
   int coeff_ctx = 0;
@@ -668,7 +673,7 @@
 
       if (level > NUM_BASE_LEVELS) {
         int ctx;
-        ctx = get_br_ctx(levels, scan[c], bwl);
+        ctx = get_br_ctx(levels, scan[c], bwl, level_counts[scan[c]]);
         int base_range = level - 1 - NUM_BASE_LEVELS;
         if (base_range < COEFF_BASE_RANGE) {
           cost += coeff_costs->lps_cost[ctx][base_range];
@@ -1461,7 +1466,13 @@
 #endif
 
     if (abs_qc > NUM_BASE_LEVELS) {
-      int ctx = get_br_ctx(txb_info->levels, scan[scan_idx], txb_info->bwl);
+      const int row = scan[scan_idx] >> txb_info->bwl;
+      const int col = scan[scan_idx] - (row << txb_info->bwl);
+      const int count = get_level_count(
+          txb_info->levels, (1 << txb_info->bwl) + TX_PAD_HOR, row, col,
+          NUM_BASE_LEVELS, br_ref_offset, BR_CONTEXT_POSITION_NUM);
+      const int ctx =
+          get_br_ctx(txb_info->levels, scan[scan_idx], txb_info->bwl, count);
       cost += get_br_cost(abs_qc, ctx, txb_costs->lps_cost[ctx]);
       cost += get_golomb_cost(abs_qc);
     }
@@ -2022,6 +2033,7 @@
   const int height = tx_size_high[tx_size];
   uint8_t levels_buf[TX_PAD_2D];
   uint8_t *const levels = set_levels(levels_buf, width);
+  uint8_t level_counts[MAX_TX_SQUARE];
   const uint8_t allow_update_cdf = args->allow_update_cdf;
 
   TX_SIZE txsize_ctx = get_txsize_context(tx_size);
@@ -2098,7 +2110,6 @@
 
 #if !USE_CAUSAL_BASE_CTX
   // Reverse process order to handle coefficient level and sign.
-  uint8_t level_counts[MAX_TX_SQUARE];
   for (int i = 0; i < NUM_BASE_LEVELS; ++i) {
     av1_get_base_level_counts(levels, i, width, height, level_counts);
     for (c = eob - 1; c >= 0; --c) {
@@ -2138,51 +2149,54 @@
     x->mbmi_ext->dc_sign_ctx[plane][block] = dc_sign_ctx;
   }
 
-  for (c = update_eob; c >= 0; --c) {
-    const tran_low_t level = abs(tcoeff[scan[c]]);
-    int idx;
-    int ctx;
+  if (update_eob >= 0) {
+    av1_get_br_level_counts(levels, width, height, level_counts);
+    for (c = update_eob; c >= 0; --c) {
+      const tran_low_t level = abs(tcoeff[scan[c]]);
+      int idx;
+      int ctx;
 
-    if (level <= NUM_BASE_LEVELS) continue;
+      if (level <= NUM_BASE_LEVELS) continue;
 
-    // level is above 1.
-    ctx = get_br_ctx(levels, scan[c], bwl);
+      // level is above 1.
+      ctx = get_br_ctx(levels, scan[c], bwl, level_counts[scan[c]]);
 
-    int base_range = level - 1 - NUM_BASE_LEVELS;
-    int br_set_idx = base_range < COEFF_BASE_RANGE
-                         ? coeff_to_br_index[base_range]
-                         : BASE_RANGE_SETS;
+      int base_range = level - 1 - NUM_BASE_LEVELS;
+      int br_set_idx = base_range < COEFF_BASE_RANGE
+                           ? coeff_to_br_index[base_range]
+                           : BASE_RANGE_SETS;
 
-    for (idx = 0; idx < BASE_RANGE_SETS; ++idx) {
-      if (idx == br_set_idx) {
-        int br_base = br_index_to_coeff[br_set_idx];
-        int br_offset = base_range - br_base;
-        ++td->counts->coeff_br[txsize_ctx][plane_type][idx][ctx][1];
-        if (allow_update_cdf)
-          update_bin(ec_ctx->coeff_br_cdf[txsize_ctx][plane_type][idx][ctx], 1,
-                     2);
-        int extra_bits = (1 << br_extra_bits[idx]) - 1;
-        for (int tok = 0; tok < extra_bits; ++tok) {
-          if (br_offset == tok) {
-            ++td->counts->coeff_lps[txsize_ctx][plane_type][ctx][1];
-            if (allow_update_cdf)
-              update_bin(ec_ctx->coeff_lps_cdf[txsize_ctx][plane_type][ctx], 1,
-                         2);
-            break;
-          }
-          ++td->counts->coeff_lps[txsize_ctx][plane_type][ctx][0];
+      for (idx = 0; idx < BASE_RANGE_SETS; ++idx) {
+        if (idx == br_set_idx) {
+          int br_base = br_index_to_coeff[br_set_idx];
+          int br_offset = base_range - br_base;
+          ++td->counts->coeff_br[txsize_ctx][plane_type][idx][ctx][1];
           if (allow_update_cdf)
-            update_bin(ec_ctx->coeff_lps_cdf[txsize_ctx][plane_type][ctx], 0,
-                       2);
+            update_bin(ec_ctx->coeff_br_cdf[txsize_ctx][plane_type][idx][ctx],
+                       1, 2);
+          int extra_bits = (1 << br_extra_bits[idx]) - 1;
+          for (int tok = 0; tok < extra_bits; ++tok) {
+            if (br_offset == tok) {
+              ++td->counts->coeff_lps[txsize_ctx][plane_type][ctx][1];
+              if (allow_update_cdf)
+                update_bin(ec_ctx->coeff_lps_cdf[txsize_ctx][plane_type][ctx],
+                           1, 2);
+              break;
+            }
+            ++td->counts->coeff_lps[txsize_ctx][plane_type][ctx][0];
+            if (allow_update_cdf)
+              update_bin(ec_ctx->coeff_lps_cdf[txsize_ctx][plane_type][ctx], 0,
+                         2);
+          }
+          break;
         }
-        break;
+        ++td->counts->coeff_br[txsize_ctx][plane_type][idx][ctx][0];
+        if (allow_update_cdf)
+          update_bin(ec_ctx->coeff_br_cdf[txsize_ctx][plane_type][idx][ctx], 0,
+                     2);
       }
-      ++td->counts->coeff_br[txsize_ctx][plane_type][idx][ctx][0];
-      if (allow_update_cdf)
-        update_bin(ec_ctx->coeff_br_cdf[txsize_ctx][plane_type][idx][ctx], 0,
-                   2);
+      // use 0-th order Golomb code to handle the residual level.
     }
-    // use 0-th order Golomb code to handle the residual level.
   }
 
   int cul_level = av1_get_txb_entropy_context(tcoeff, scan_order, eob);