Clean av1_read_coeffs_txb()

Let read_nz_map() etc. output to uint8_t*.

Temporarily introduced coeff_is_byte_flag as an argument of
get_nz_count() and get_nz_map_ctx() to handle different types of input
coefficients. It helps to not duplicate too many functions.
If possible, unify types to uint8_t* later.

Change-Id: Idbbe67b7ec563f8c9299daa0fa7d05e15c8295e6
diff --git a/av1/common/txb_common.h b/av1/common/txb_common.h
index 0aa34f8..0fc93f9 100644
--- a/av1/common/txb_common.h
+++ b/av1/common/txb_common.h
@@ -39,6 +39,9 @@
   /* clang-format on*/
 };
 
+// TODO(linfengz): Some functions have coeff_is_byte_flag to handle different
+// types of input coefficients. If possible, unify types to uint8_t* later.
+
 static INLINE int get_level_count(const tran_low_t *tcoeffs, int bwl,
                                   int height, int row, int col, int level,
                                   int (*nb_offset)[2], int nb_num) {
@@ -306,8 +309,10 @@
 };
 
 #if REDUCE_CONTEXT_DEPENDENCY
-static INLINE int get_nz_count(const tran_low_t *tcoeffs, int bwl, int height,
-                               int row, int col, int prev_row, int prev_col) {
+static INLINE int get_nz_count(const void *const tcoeffs, const int bwl,
+                               const int height, const int row, const int col,
+                               const int coeff_is_byte_flag, const int prev_row,
+                               const int prev_col) {
   int count = 0;
   for (int idx = 0; idx < SIG_REF_OFFSET_NUM; ++idx) {
     const int ref_row = row + sig_ref_offset[idx][0];
@@ -316,13 +321,16 @@
         ref_col >= (1 << bwl) || (prev_row == ref_row && prev_col == ref_col))
       continue;
     const int nb_pos = (ref_row << bwl) + ref_col;
-    count += (tcoeffs[nb_pos] != 0);
+    count +=
+        ((coeff_is_byte_flag ? ((const uint8_t *)tcoeffs)[nb_pos]
+                             : ((const tran_low_t *)tcoeffs)[nb_pos]) != 0);
   }
   return count;
 }
 #else
-static INLINE int get_nz_count(const tran_low_t *tcoeffs, int bwl, int height,
-                               int row, int col) {
+static INLINE int get_nz_count(const void *const tcoeffs, const int bwl,
+                               const int height, const int row, const int col,
+                               const int coeff_is_byte_flag) {
   int count = 0;
   for (int idx = 0; idx < SIG_REF_OFFSET_NUM; ++idx) {
     const int ref_row = row + sig_ref_offset[idx][0];
@@ -331,7 +339,9 @@
         ref_col >= (1 << bwl))
       continue;
     const int nb_pos = (ref_row << bwl) + ref_col;
-    count += (tcoeffs[nb_pos] != 0);
+    count +=
+        ((coeff_is_byte_flag ? ((const uint8_t *)tcoeffs)[nb_pos]
+                             : ((const tran_low_t *)tcoeffs)[nb_pos]) != 0);
   }
   return count;
 }
@@ -408,9 +418,10 @@
   return offset + 12 + ctx;
 }
 
-static INLINE int get_nz_map_ctx(const tran_low_t *tcoeffs, const int scan_idx,
-                                 const int16_t *scan, const int bwl,
-                                 const int height, TX_TYPE tx_type) {
+static INLINE int get_nz_map_ctx(const void *const tcoeffs, const int scan_idx,
+                                 const int16_t *const scan, const int bwl,
+                                 const int height, const TX_TYPE tx_type,
+                                 const int coeff_is_byte_flag) {
   const int coeff_idx = scan[scan_idx];
   const int row = coeff_idx >> bwl;
   const int col = coeff_idx - (row << bwl);
@@ -427,9 +438,11 @@
     prev_row = -1;
     prev_col = -1;
   }
-  int count = get_nz_count(tcoeffs, bwl, height, row, col, prev_row, prev_col);
+  const int count = get_nz_count(tcoeffs, bwl, height, row, col,
+                                 coeff_is_byte_flag, prev_row, prev_col);
 #else
-  int count = get_nz_count(tcoeffs, bwl, height, row, col);
+  const int count =
+      get_nz_count(tcoeffs, bwl, height, row, col, coeff_is_byte_flag);
 #endif
   return get_nz_map_ctx_from_count(count, coeff_idx, bwl, tx_type);
 }
diff --git a/av1/decoder/decodetxb.c b/av1/decoder/decodetxb.c
index 27f8148..e5622e2 100644
--- a/av1/decoder/decodetxb.c
+++ b/av1/decoder/decodetxb.c
@@ -45,11 +45,12 @@
   return x - 1;
 }
 
-static INLINE int read_nz_map(aom_reader *r, tran_low_t *tcoeffs, int plane,
-                              const int16_t *scan, TX_SIZE tx_size,
-                              TX_TYPE tx_type, FRAME_CONTEXT *fc,
-                              FRAME_COUNTS *counts) {
-  TX_SIZE txs_ctx = get_txsize_context(tx_size);
+static INLINE int read_nz_map(aom_reader *const r, uint8_t *const levels,
+                              const int plane, const int16_t *const scan,
+                              const TX_SIZE tx_size, const TX_TYPE tx_type,
+                              FRAME_CONTEXT *const fc,
+                              FRAME_COUNTS *const counts) {
+  const TX_SIZE txs_ctx = get_txsize_context(tx_size);
   const int bwl = b_width_log2_lookup[txsize_to_bsize[tx_size]] + 2;
   const int height = tx_size_high[tx_size];
 #if CONFIG_CTX1D
@@ -71,7 +72,7 @@
   int c;
   for (c = 0; c < seg_eob; ++c) {
     int is_nz;
-    int coeff_ctx = get_nz_map_ctx(tcoeffs, c, scan, bwl, height, tx_type);
+    int coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type, 1);
     int eob_ctx = get_eob_ctx(scan[c], txs_ctx, tx_type);
 
     if (c < seg_eob - 1) {
@@ -87,7 +88,7 @@
     }
 
     // set non-zero coefficient map.
-    tcoeffs[scan[c]] = is_nz;
+    levels[scan[c]] = is_nz;
 
     if (c == seg_eob - 1) {
       ++c;
@@ -112,11 +113,10 @@
 }
 
 #if CONFIG_CTX1D
-static INLINE int read_nz_map_vert(aom_reader *r, tran_low_t *tcoeffs,
-                                   int plane, const int16_t *scan,
-                                   const int16_t *iscan, TX_SIZE tx_size,
-                                   TX_TYPE tx_type, FRAME_CONTEXT *fc,
-                                   FRAME_COUNTS *counts) {
+static INLINE int read_nz_map_vert(aom_reader *r, uint8_t *levels, int plane,
+                                   const int16_t *scan, const int16_t *iscan,
+                                   TX_SIZE tx_size, TX_TYPE tx_type,
+                                   FRAME_CONTEXT *fc, FRAME_COUNTS *counts) {
   const TX_SIZE txs_ctx = get_txsize_context(tx_size);
   const PLANE_TYPE plane_type = get_plane_type(plane);
   const TX_CLASS tx_class = get_tx_class(tx_type);
@@ -147,7 +147,7 @@
           int coeff_idx = row * width + col;
           int scan_idx = iscan[coeff_idx];
           int coeff_ctx =
-              get_nz_map_ctx(tcoeffs, scan_idx, scan, bwl, height, tx_type);
+              get_nz_map_ctx(levels, scan_idx, scan, bwl, height, tx_type, 1);
 #if LV_MAP_PROB
           int is_nz = av1_read_record_bin(
               counts, r, fc->nz_map_cdf[txs_ctx][plane_type][coeff_ctx], 2,
@@ -156,7 +156,7 @@
           int is_nz = aom_read(r, nz_map[coeff_ctx], ACCT_STR);
 #endif
           if (counts) ++counts->nz_map[txs_ctx][plane_type][coeff_ctx][is_nz];
-          tcoeffs[coeff_idx] = is_nz;
+          levels[coeff_idx] = is_nz;
           if (is_nz) {
             eob = AOMMAX(eob, iscan[coeff_idx] + 1);
             if (row + 1 != height) {
@@ -179,7 +179,7 @@
           }
         } else {
           int coeff_idx = row * width + col;
-          tcoeffs[coeff_idx] = 1;
+          levels[coeff_idx] = 1;
           eob = AOMMAX(eob, iscan[coeff_idx] + 1);
         }
       }
@@ -191,11 +191,10 @@
   return eob;
 }
 
-static INLINE int read_nz_map_horiz(aom_reader *r, tran_low_t *tcoeffs,
-                                    int plane, const int16_t *scan,
-                                    const int16_t *iscan, TX_SIZE tx_size,
-                                    TX_TYPE tx_type, FRAME_CONTEXT *fc,
-                                    FRAME_COUNTS *counts) {
+static INLINE int read_nz_map_horiz(aom_reader *r, uint8_t *levels, int plane,
+                                    const int16_t *scan, const int16_t *iscan,
+                                    TX_SIZE tx_size, TX_TYPE tx_type,
+                                    FRAME_CONTEXT *fc, FRAME_COUNTS *counts) {
   const TX_SIZE txs_ctx = get_txsize_context(tx_size);
   const PLANE_TYPE plane_type = get_plane_type(plane);
   const TX_CLASS tx_class = get_tx_class(tx_type);
@@ -226,7 +225,7 @@
           int coeff_idx = row * width + col;
           int scan_idx = iscan[coeff_idx];
           int coeff_ctx =
-              get_nz_map_ctx(tcoeffs, scan_idx, scan, bwl, height, tx_type);
+              get_nz_map_ctx(levels, scan_idx, scan, bwl, height, tx_type, 1);
 #if LV_MAP_PROB
           int is_nz = av1_read_record_bin(
               counts, r, fc->nz_map_cdf[txs_ctx][plane_type][coeff_ctx], 2,
@@ -235,7 +234,7 @@
           int is_nz = aom_read(r, nz_map[coeff_ctx], ACCT_STR);
 #endif
           if (counts) ++counts->nz_map[txs_ctx][plane_type][coeff_ctx][is_nz];
-          tcoeffs[coeff_idx] = is_nz;
+          levels[coeff_idx] = is_nz;
           if (is_nz) {
             eob = AOMMAX(eob, iscan[coeff_idx] + 1);
             int eob_ctx = get_hv_eob_ctx(row, col, eob_ls);
@@ -255,7 +254,7 @@
           }
         } else {
           int coeff_idx = row * width + col;
-          tcoeffs[coeff_idx] = 1;
+          levels[coeff_idx] = 1;
           eob = AOMMAX(eob, iscan[coeff_idx] + 1);
         }
       }
@@ -314,6 +313,7 @@
     return 0;
   }
 
+  memset(levels, 0, sizeof(levels[0]) * seg_eob);
   memset(signs, 0, sizeof(signs[0]) * seg_eob);
 
   (void)blk_row;
@@ -332,7 +332,7 @@
   TX_CLASS tx_class = get_tx_class(tx_type);
   if (tx_class == TX_CLASS_2D) {
     *eob =
-        read_nz_map(r, tcoeffs, plane, scan, tx_size, tx_type, ec_ctx, counts);
+        read_nz_map(r, levels, plane, scan, tx_size, tx_type, ec_ctx, counts);
   } else {
 #if LV_MAP_PROB
     const int eob_mode = av1_read_record_bin(
@@ -344,28 +344,24 @@
 #endif
     if (counts) ++counts->eob_mode[txs_ctx][plane_type][tx_class][eob_mode];
     if (eob_mode == 0) {
-      *eob = read_nz_map(r, tcoeffs, plane, scan, tx_size, tx_type, ec_ctx,
-                         counts);
+      *eob =
+          read_nz_map(r, levels, plane, scan, tx_size, tx_type, ec_ctx, counts);
     } else {
       assert(tx_class == TX_CLASS_VERT || tx_class == TX_CLASS_HORIZ);
       if (tx_class == TX_CLASS_VERT)
-        *eob = read_nz_map_vert(r, tcoeffs, plane, scan, iscan, tx_size,
-                                tx_type, ec_ctx, counts);
+        *eob = read_nz_map_vert(r, levels, plane, scan, iscan, tx_size, tx_type,
+                                ec_ctx, counts);
       else
-        *eob = read_nz_map_horiz(r, tcoeffs, plane, scan, iscan, tx_size,
+        *eob = read_nz_map_horiz(r, levels, plane, scan, iscan, tx_size,
                                  tx_type, ec_ctx, counts);
     }
   }
 #else
-  *eob = read_nz_map(r, tcoeffs, plane, scan, tx_size, tx_type, ec_ctx, counts);
+  *eob = read_nz_map(r, levels, plane, scan, tx_size, tx_type, ec_ctx, counts);
 #endif
   *max_scan_line = *eob;
 
   int i;
-  for (i = 0; i < seg_eob; i++) {
-    levels[i] = (uint8_t)tcoeffs[i];
-  }
-
   for (i = 0; i < NUM_BASE_LEVELS; ++i) {
 #if !LV_MAP_PROB
     aom_prob *coeff_base = ec_ctx->coeff_base[txs_ctx][plane_type][i];
diff --git a/av1/encoder/encodetxb.c b/av1/encoder/encodetxb.c
index bbb20d8..f53c76e 100644
--- a/av1/encoder/encodetxb.c
+++ b/av1/encoder/encodetxb.c
@@ -116,7 +116,7 @@
 #endif
 
   for (int c = 0; c < eob; ++c) {
-    int coeff_ctx = get_nz_map_ctx(tcoeff, c, scan, bwl, height, tx_type);
+    int coeff_ctx = get_nz_map_ctx(tcoeff, c, scan, bwl, height, tx_type, 0);
     int eob_ctx = get_eob_ctx(scan[c], txs_ctx, tx_type);
 
     tran_low_t v = tcoeff[scan[c]];
@@ -177,7 +177,7 @@
           int scan_idx = iscan[coeff_idx];
           int is_nz = tcoeff[coeff_idx] != 0;
           int coeff_ctx =
-              get_nz_map_ctx(tcoeff, scan_idx, scan, bwl, height, tx_type);
+              get_nz_map_ctx(tcoeff, scan_idx, scan, bwl, height, tx_type, 0);
 #if LV_MAP_PROB
           aom_write_bin(w, is_nz,
                         fc->nz_map_cdf[txs_ctx][plane_type][coeff_ctx], 2);
@@ -236,7 +236,7 @@
           int scan_idx = iscan[coeff_idx];
           int is_nz = tcoeff[coeff_idx] != 0;
           int coeff_ctx =
-              get_nz_map_ctx(tcoeff, scan_idx, scan, bwl, height, tx_type);
+              get_nz_map_ctx(tcoeff, scan_idx, scan, bwl, height, tx_type, 0);
 #if LV_MAP_PROB
           aom_write_bin(w, is_nz,
                         fc->nz_map_cdf[txs_ctx][plane_type][coeff_ctx], 2);
@@ -583,7 +583,7 @@
     tran_low_t v = qcoeff[scan[c]];
     int is_nz = (v != 0);
     if (c + 1 != seg_eob) {
-      int coeff_ctx = get_nz_map_ctx(qcoeff, c, scan, bwl, height, tx_type);
+      int coeff_ctx = get_nz_map_ctx(qcoeff, c, scan, bwl, height, tx_type, 0);
       cost += coeff_costs->nz_map_cost[coeff_ctx][is_nz];
       if (is_nz) {
         int eob_ctx = get_eob_ctx(scan[c], txs_ctx, tx_type);
@@ -624,7 +624,7 @@
           int scan_idx = iscan[coeff_idx];
           int is_nz = qcoeff[coeff_idx] != 0;
           int coeff_ctx =
-              get_nz_map_ctx(qcoeff, scan_idx, scan, bwl, height, tx_type);
+              get_nz_map_ctx(qcoeff, scan_idx, scan, bwl, height, tx_type, 0);
           cost += coeff_costs->nz_map_cost[coeff_ctx][is_nz];
           if (is_nz) {
             int eob_ctx = get_hv_eob_ctx(c, r, eob_ls);
@@ -664,7 +664,7 @@
           int scan_idx = iscan[coeff_idx];
           int is_nz = qcoeff[coeff_idx] != 0;
           int coeff_ctx =
-              get_nz_map_ctx(qcoeff, scan_idx, scan, bwl, height, tx_type);
+              get_nz_map_ctx(qcoeff, scan_idx, scan, bwl, height, tx_type, 0);
           cost += coeff_costs->nz_map_cost[coeff_ctx][is_nz];
           if (is_nz) {
             int eob_ctx = get_hv_eob_ctx(r, c, eob_ls);
@@ -858,7 +858,7 @@
   const int16_t *scan = txb_info->scan_order->scan;
   const int bwl = txb_info->bwl;
   const int height = txb_info->height;
-  tran_low_t *qcoeff = txb_info->qcoeff;
+  const tran_low_t *const qcoeff = txb_info->qcoeff;
   const BASE_CTX_TABLE *base_ctx_table =
       txb_info->coeff_ctx_table->base_ctx_table;
   for (int c = 0; c < txb_info->eob; ++c) {
@@ -879,10 +879,10 @@
       prev_col = -1;
     }
     txb_cache->nz_count_arr[coeff_idx] =
-        get_nz_count(qcoeff, bwl, height, row, col, prev_row, prev_col);
+        get_nz_count(qcoeff, bwl, height, row, col, 0, prev_row, prev_col);
 #else
     txb_cache->nz_count_arr[coeff_idx] =
-        get_nz_count(qcoeff, bwl, height, row, col);
+        get_nz_count(qcoeff, bwl, height, row, col, 0);
 #endif
     const int nz_count = txb_cache->nz_count_arr[coeff_idx];
     txb_cache->nz_ctx_arr[coeff_idx] =
@@ -1487,7 +1487,7 @@
         txb_cache->nz_ctx_arr[nb_coeff_idx] = get_nz_map_ctx_from_count(
             count, nb_coeff_idx, txb_info->bwl, txb_info->tx_type);
         // int ref_ctx = get_nz_map_ctx(txb_info->qcoeff, nb_coeff_idx,
-        // txb_info->bwl, tx_type);
+        // txb_info->bwl, tx_type, 0);
         // 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);
@@ -1579,7 +1579,7 @@
   if (scan_idx < txb_info->seg_eob) {
     int coeff_ctx =
         get_nz_map_ctx(txb_info->qcoeff, scan_idx, scan, txb_info->bwl,
-                       txb_info->height, txb_info->tx_type);
+                       txb_info->height, txb_info->tx_type, 0);
     cost += txb_costs->nz_map_cost[coeff_ctx][is_nz];
   }
 
@@ -2027,7 +2027,7 @@
   for (int c = 0; c < eob; ++c) {
     tran_low_t v = tcoeff[scan[c]];
     int is_nz = (v != 0);
-    int coeff_ctx = get_nz_map_ctx(tcoeff, c, scan, bwl, height, tx_type);
+    int coeff_ctx = get_nz_map_ctx(tcoeff, c, scan, bwl, height, tx_type, 0);
     int eob_ctx = get_eob_ctx(scan[c], txsize_ctx, tx_type);
 
     if (c == seg_eob - 1) break;
@@ -2079,7 +2079,7 @@
           int scan_idx = iscan[coeff_idx];
           int is_nz = tcoeff[coeff_idx] != 0;
           int coeff_ctx =
-              get_nz_map_ctx(tcoeff, scan_idx, scan, bwl, height, tx_type);
+              get_nz_map_ctx(tcoeff, scan_idx, scan, bwl, height, tx_type, 0);
           ++(*nz_map_count)[coeff_ctx][is_nz];
 #if LV_MAP_PROB
           update_bin(fc->nz_map_cdf[txs_ctx][plane_type][coeff_ctx], is_nz, 2);
@@ -2130,7 +2130,7 @@
           int scan_idx = iscan[coeff_idx];
           int is_nz = tcoeff[coeff_idx] != 0;
           int coeff_ctx =
-              get_nz_map_ctx(tcoeff, scan_idx, scan, bwl, height, tx_type);
+              get_nz_map_ctx(tcoeff, scan_idx, scan, bwl, height, tx_type, 0);
           ++(*nz_map_count)[coeff_ctx][is_nz];
 #if LV_MAP_PROB
           update_bin(fc->nz_map_cdf[txs_ctx][plane_type][coeff_ctx], is_nz, 2);