Add levels array in av1_optimize_txb()

av1_txb_init_levels() has prefix 'av1' because it will be optimized
and not static later.

Change-Id: I988da817335f122522a76c4412207c3a6fdd6b71
diff --git a/av1/encoder/encodetxb.c b/av1/encoder/encodetxb.c
index 05a2033..4c3df5f 100644
--- a/av1/encoder/encodetxb.c
+++ b/av1/encoder/encodetxb.c
@@ -189,8 +189,9 @@
   return eob_cost;
 }
 
-extern int get_coeff_cost(tran_low_t qc, int scan_idx, TxbInfo *txb_info,
-                          const LV_MAP_COEFF_COST *txb_costs);
+static int get_coeff_cost(const tran_low_t qc, const int scan_idx,
+                          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,
@@ -198,6 +199,7 @@
   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];
+  const uint8_t *const levels = txb_info->levels;
   stats->new_eob = -1;
   stats->update = 0;
 
@@ -230,9 +232,8 @@
   stats->rate_low = low_qc_cost;
   stats->rd_low = RDCOST(txb_info->rdmult, stats->rate_low, stats->dist_low);
 
-  int coeff_ctx =
-      get_nz_map_ctx(txb_info->qcoeff, scan_idx, scan, txb_info->bwl,
-                     txb_info->height, txb_info->tx_type, 0);
+  int coeff_ctx = get_nz_map_ctx(levels, scan_idx, scan, txb_info->bwl,
+                                 txb_info->height, txb_info->tx_type, 1);
 
   if ((stats->rd_low < stats->rd) && (stats->low_qc == 0)) {
     stats->nz_rate = txb_costs->nz_map_cost[coeff_ctx][0];
@@ -241,12 +242,26 @@
   }
 }
 
-static void update_coeff(int coeff_idx, tran_low_t qc, TxbInfo *txb_info) {
+static INLINE void update_qcoeff(const int coeff_idx, const tran_low_t qc,
+                                 const TxbInfo *const txb_info) {
   txb_info->qcoeff[coeff_idx] = qc;
+  txb_info->levels[coeff_idx] = (uint8_t)clamp(abs(qc), 0, UINT8_MAX);
+}
+
+static INLINE void update_coeff(const int coeff_idx, const tran_low_t qc,
+                                const TxbInfo *const txb_info) {
+  update_qcoeff(coeff_idx, qc, txb_info);
   const int dqv = txb_info->dequant[coeff_idx != 0];
   txb_info->dqcoeff[coeff_idx] = qcoeff_to_dqcoeff(qc, dqv, txb_info->shift);
 }
 
+static INLINE void av1_txb_init_levels(const tran_low_t *const coeff,
+                                       const int size, uint8_t *const levels) {
+  for (int i = 0; i < size; i++) {
+    levels[i] = (uint8_t)clamp(abs(coeff[i]), 0, UINT8_MAX);
+  }
+}
+
 void av1_write_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
                           aom_writer *w, int blk_row, int blk_col, int block,
                           int plane, TX_SIZE tx_size, const tran_low_t *tcoeff,
@@ -264,8 +279,7 @@
   const int height = tx_size_high[tx_size];
   uint16_t update_eob = 0;
   FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
-  uint8_t levels[64 * 64];
-  int8_t signs[64 * 64];
+  uint8_t levels[MAX_TX_SQUARE];
 
   (void)blk_row;
   (void)blk_col;
@@ -275,10 +289,7 @@
 
   if (eob == 0) return;
 
-  for (int i = 0; i < seg_eob; i++) {
-    levels[i] = (uint8_t)clamp(abs(tcoeff[i]), 0, UINT8_MAX);
-    signs[i] = (int8_t)(tcoeff[i] < 0);
-  }
+  av1_txb_init_levels(tcoeff, seg_eob, levels);
 
 #if CONFIG_TXK_SEL
   av1_write_tx_type(cm, xd, blk_row, blk_col, block, plane,
@@ -345,7 +356,7 @@
 #else
   for (int i = 1; i < eob; ++i) {
     c = eob - 1 - i;
-    int coeff_ctx = get_nz_map_ctx(tcoeff, c, scan, bwl, height, tx_type, 0);
+    int coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type, 1);
 
     tran_low_t v = tcoeff[scan[c]];
     int is_nz = (v != 0);
@@ -357,7 +368,7 @@
   for (int i = 0; i < NUM_BASE_LEVELS; ++i) {
     update_eob = 0;
     for (c = eob - 1; c >= 0; --c) {
-      const int level = levels[scan[c]];
+      const tran_low_t level = abs(tcoeff[scan[c]]);
       int ctx;
 
       if (level <= i) continue;
@@ -380,8 +391,9 @@
   // Loop to code all signs in the transform block,
   // starting with the sign of DC (if applicable)
   for (c = 0; c < eob; ++c) {
-    const int level = levels[scan[c]];
-    const int sign = signs[scan[c]];
+    const tran_low_t v = tcoeff[scan[c]];
+    const tran_low_t level = abs(v);
+    const int sign = (v < 0) ? 1 : 0;
     if (level == 0) continue;
 
     if (c == 0) {
@@ -397,7 +409,7 @@
   }
 
   for (c = update_eob; c >= 0; --c) {
-    const int level = levels[scan[c]];
+    const tran_low_t level = abs(tcoeff[scan[c]]);
     int idx;
     int ctx;
 
@@ -561,6 +573,7 @@
 
   const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, mbmi);
   const int16_t *scan = scan_order->scan;
+  uint8_t levels[MAX_TX_SQUARE];
 
   LV_MAP_COEFF_COST *coeff_costs = &x->coeff_costs[txs_ctx][plane_type];
 
@@ -572,6 +585,8 @@
   }
   cost = coeff_costs->txb_skip_cost[txb_skip_ctx][0];
 
+  av1_txb_init_levels(qcoeff, tx_size_2d[tx_size], levels);
+
 #if CONFIG_TXK_SEL
   cost += av1_tx_type_cost(cm, x, xd, mbmi->sb_type, plane, tx_size, tx_type);
 #endif
@@ -593,7 +608,7 @@
 
     if (c < eob - 1) {
 #if !USE_CAUSAL_BASE_CTX
-      int coeff_ctx = get_nz_map_ctx(qcoeff, c, scan, bwl, height, tx_type, 0);
+      int coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type, 1);
 #endif
       cost += coeff_costs->nz_map_cost[coeff_ctx][is_nz];
     }
@@ -925,14 +940,14 @@
   if (nb_scan_idx < scan_idx) {
     const int count = txb_cache->nz_count_arr[coeff_idx];
     assert(count > 0);
-    txb_info->qcoeff[nb_coeff_idx] = get_lower_coeff(nb_coeff);
+    update_qcoeff(nb_coeff_idx, get_lower_coeff(nb_coeff), txb_info);
     const int new_ctx = get_nz_map_ctx_from_count(
         count - 1, coeff_idx, txb_info->bwl, txb_info->height,
 #if USE_CAUSAL_BASE_CTX
         0,
 #endif
         txb_info->tx_type);
-    txb_info->qcoeff[nb_coeff_idx] = nb_coeff;
+    update_qcoeff(nb_coeff_idx, nb_coeff, txb_info);
     const int ctx = txb_cache->nz_ctx_arr[coeff_idx];
     const int is_nz = abs_qc > 0;
     const int org_cost = txb_costs->nz_map_cost[ctx][is_nz];
@@ -1261,15 +1276,13 @@
 }
 
 // TODO(angiebird): add static to this function it's called
-void update_level_down(int coeff_idx, TxbCache *txb_cache, TxbInfo *txb_info) {
+void update_level_down(const int coeff_idx, TxbCache *const txb_cache,
+                       TxbInfo *const txb_info) {
   const tran_low_t qc = txb_info->qcoeff[coeff_idx];
   const int abs_qc = abs(qc);
   if (qc == 0) return;
   const tran_low_t low_coeff = get_lower_coeff(qc);
-  txb_info->qcoeff[coeff_idx] = low_coeff;
-  const int dqv = txb_info->dequant[coeff_idx != 0];
-  txb_info->dqcoeff[coeff_idx] =
-      qcoeff_to_dqcoeff(low_coeff, dqv, txb_info->shift);
+  update_coeff(coeff_idx, low_coeff, txb_info);
 
   const int row = coeff_idx >> txb_info->bwl;
   const int col = coeff_idx - (row << txb_info->bwl);
@@ -1300,8 +1313,8 @@
             0,
 #endif
             txb_info->tx_type);
-        // int ref_ctx = get_nz_map_ctx(txb_info->qcoeff, nb_coeff_idx,
-        // txb_info->bwl, tx_type, 0);
+        // int ref_ctx = get_nz_map_ctx(txb_info->levels, nb_coeff_idx,
+        // txb_info->bwl, tx_type, 1);
         // if (ref_ctx != txb_cache->nz_ctx_arr[nb_coeff_idx])
         //   printf("nz ctx %d ref_ctx %d\n",
         //   txb_cache->nz_ctx_arr[nb_coeff_idx], ref_ctx);
@@ -1382,8 +1395,9 @@
   }
 }
 
-int get_coeff_cost(tran_low_t qc, int scan_idx, TxbInfo *txb_info,
-                   const LV_MAP_COEFF_COST *txb_costs) {
+static int get_coeff_cost(const tran_low_t qc, const int scan_idx,
+                          const TxbInfo *const txb_info,
+                          const LV_MAP_COEFF_COST *const txb_costs) {
   const TXB_CTX *txb_ctx = txb_info->txb_ctx;
   const int is_nz = (qc != 0);
   const tran_low_t abs_qc = abs(qc);
@@ -1399,8 +1413,8 @@
   if (scan_idx < txb_info->eob - 1) {
 #if !USE_CAUSAL_BASE_CTX
     int coeff_ctx =
-        get_nz_map_ctx(txb_info->qcoeff, scan_idx, scan, txb_info->bwl,
-                       txb_info->height, txb_info->tx_type, 0);
+        get_nz_map_ctx(txb_info->levels, scan_idx, scan, txb_info->bwl,
+                       txb_info->height, txb_info->tx_type, 1);
 #endif
     cost += txb_costs->nz_map_cost[coeff_ctx][is_nz];
   }
@@ -1474,7 +1488,7 @@
       org_cost += cost;
     }
   }
-  txb_info->qcoeff[coeff_idx] = get_lower_coeff(qc);
+  update_qcoeff(coeff_idx, get_lower_coeff(qc), txb_info);
   int new_cost = 0;
   for (int i = 0; i < ALL_REF_OFFSET_NUM; ++i) {
     int nb_row = row - all_ref_offset[i][0];
@@ -1491,7 +1505,7 @@
       new_cost += cost;
     }
   }
-  txb_info->qcoeff[coeff_idx] = qc;
+  update_qcoeff(coeff_idx, qc, txb_info);
   return new_cost - org_cost;
 }
 
@@ -1551,11 +1565,11 @@
       int org_eob = txb_info->eob;
       int cost = get_txb_cost(txb_info, txb_costs);
 
-      txb_info->qcoeff[last_ci] = get_lower_coeff(last_coeff);
+      update_qcoeff(last_ci, get_lower_coeff(last_coeff), txb_info);
       set_eob(txb_info, new_eob);
       int new_cost = get_txb_cost(txb_info, txb_costs);
       set_eob(txb_info, org_eob);
-      txb_info->qcoeff[last_ci] = last_coeff;
+      update_qcoeff(last_ci, last_coeff, txb_info);
 
       int ref_cost_diff = -cost + new_cost;
       if (cost_diff != ref_cost_diff)
@@ -1632,15 +1646,19 @@
 
   tran_low_t *org_qcoeff = txb_info->qcoeff;
   tran_low_t *org_dqcoeff = txb_info->dqcoeff;
+  uint8_t *const org_levels = txb_info->levels;
 
   tran_low_t tmp_qcoeff[MAX_TX_SQUARE];
   tran_low_t tmp_dqcoeff[MAX_TX_SQUARE];
+  uint8_t tmp_levels[MAX_TX_SQUARE];
   const int org_eob = txb_info->eob;
   if (dry_run) {
     memcpy(tmp_qcoeff, org_qcoeff, sizeof(org_qcoeff[0]) * max_eob);
     memcpy(tmp_dqcoeff, org_dqcoeff, sizeof(org_dqcoeff[0]) * max_eob);
+    memcpy(tmp_levels, org_levels, sizeof(org_levels[0]) * max_eob);
     txb_info->qcoeff = tmp_qcoeff;
     txb_info->dqcoeff = tmp_dqcoeff;
+    txb_info->levels = tmp_levels;
   }
 
   const int16_t *scan = txb_info->scan_order->scan;
@@ -1741,6 +1759,7 @@
   if (dry_run) {
     txb_info->qcoeff = org_qcoeff;
     txb_info->dqcoeff = org_dqcoeff;
+    txb_info->levels = org_levels;
     set_eob(txb_info, org_eob);
   }
   return update;
@@ -1766,15 +1785,19 @@
 
   tran_low_t *org_qcoeff = txb_info->qcoeff;
   tran_low_t *org_dqcoeff = txb_info->dqcoeff;
+  uint8_t *const org_levels = txb_info->levels;
 
   tran_low_t tmp_qcoeff[MAX_TX_SQUARE];
   tran_low_t tmp_dqcoeff[MAX_TX_SQUARE];
+  uint8_t tmp_levels[MAX_TX_SQUARE];
   const int org_eob = txb_info->eob;
   if (dry_run) {
     memcpy(tmp_qcoeff, org_qcoeff, sizeof(org_qcoeff[0]) * max_eob);
     memcpy(tmp_dqcoeff, org_dqcoeff, sizeof(org_dqcoeff[0]) * max_eob);
+    memcpy(tmp_levels, org_levels, sizeof(org_levels[0]) * max_eob);
     txb_info->qcoeff = tmp_qcoeff;
     txb_info->dqcoeff = tmp_dqcoeff;
+    txb_info->levels = tmp_levels;
   }
 
   const int16_t *scan = txb_info->scan_order->scan;
@@ -1844,6 +1867,7 @@
   if (dry_run) {
     txb_info->qcoeff = org_qcoeff;
     txb_info->dqcoeff = org_dqcoeff;
+    txb_info->levels = org_levels;
     set_eob(txb_info, org_eob);
   }
   return update;
@@ -1882,24 +1906,15 @@
   const int shift = av1_get_tx_scale(tx_size);
   const int64_t rdmult =
       (x->rdmult * plane_rd_mult[is_inter][plane_type] + 2) >> 2;
+  uint8_t levels[MAX_TX_SQUARE];
 
-  TxbInfo txb_info = { qcoeff,
-                       dqcoeff,
-                       tcoeff,
-                       dequant,
-                       shift,
-                       tx_size,
-                       txs_ctx,
-                       tx_type,
-                       bwl,
-                       stride,
-                       height,
-                       eob,
-                       seg_eob,
-                       scan_order,
-                       txb_ctx,
-                       rdmult,
-                       &cm->coeff_ctx_table };
+  TxbInfo txb_info = {
+    qcoeff,  levels,  dqcoeff,    tcoeff,  dequant, shift,
+    tx_size, txs_ctx, tx_type,    bwl,     stride,  height,
+    eob,     seg_eob, scan_order, txb_ctx, rdmult,  &cm->coeff_ctx_table
+  };
+
+  av1_txb_init_levels(qcoeff, tx_size_2d[tx_size], levels);
 
   const int update = optimize_txb(&txb_info, &txb_costs, NULL, 0, fast_mode);
 
@@ -1969,14 +1984,13 @@
   const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, mbmi);
   const int16_t *scan = scan_order->scan;
   const int seg_eob = av1_get_tx_eob(&cpi->common.seg, segment_id, tx_size);
-  int c, i;
+  int c;
   TXB_CTX txb_ctx;
   get_txb_ctx(plane_bsize, tx_size, plane, pd->above_context + blk_col,
               pd->left_context + blk_row, &txb_ctx);
   const int bwl = b_width_log2_lookup[txsize_to_bsize[tx_size]] + 2;
   const int height = tx_size_high[tx_size];
-  uint8_t levels[64 * 64];
-  int8_t signs[64 * 64];
+  uint8_t levels[MAX_TX_SQUARE];
   const uint8_t allow_update_cdf = args->allow_update_cdf;
 
   TX_SIZE txsize_ctx = get_txsize_context(tx_size);
@@ -1997,10 +2011,7 @@
     return;
   }
 
-  for (i = 0; i < seg_eob; i++) {
-    levels[i] = (uint8_t)clamp(abs(tcoeff[i]), 0, UINT8_MAX);
-    signs[i] = (int8_t)(tcoeff[i] < 0);
-  }
+  av1_txb_init_levels(tcoeff, tx_size_2d[tx_size], levels);
 
 #if CONFIG_TXK_SEL
   av1_update_tx_type_count(cm, xd, blk_row, blk_col, block, plane,
@@ -2043,7 +2054,7 @@
       }
     }
 #else
-    int coeff_ctx = get_nz_map_ctx(tcoeff, c, scan, bwl, height, tx_type, 0);
+    int coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type, 1);
 
     if (c == eob - 1) continue;
 
@@ -2056,10 +2067,10 @@
 
 #if !USE_CAUSAL_BASE_CTX
   // Reverse process order to handle coefficient level and sign.
-  for (i = 0; i < NUM_BASE_LEVELS; ++i) {
+  for (int i = 0; i < NUM_BASE_LEVELS; ++i) {
     update_eob = 0;
     for (c = eob - 1; c >= 0; --c) {
-      const int level = levels[scan[c]];
+      const tran_low_t level = abs(tcoeff[scan[c]]);
       int ctx;
 
       if (level <= i) continue;
@@ -2083,8 +2094,8 @@
 #endif
 
   c = 0;  // Update the context needed to code the DC sign (if applicable)
-  const int sign = signs[scan[c]];
-  if (levels[scan[c]] != 0) {
+  const int sign = (tcoeff[scan[c]] < 0) ? 1 : 0;
+  if (tcoeff[scan[c]] != 0) {
     int dc_sign_ctx = txb_ctx.dc_sign_ctx;
 
     ++td->counts->dc_sign[plane_type][dc_sign_ctx][sign];
@@ -2096,7 +2107,7 @@
   }
 
   for (c = update_eob; c >= 0; --c) {
-    const int level = levels[scan[c]];
+    const tran_low_t level = abs(tcoeff[scan[c]]);
     int idx;
     int ctx;