Remove temporary coeff_is_byte_flag in txb code

Change-Id: I8b3e028b3c17660b390ce9c65b3d31aabc83efd9

av1/common/txb_common.h

@@ -45,9 +45,6 @@ static const int base_ref_offset[BASE_CONTEXT_POSITION_NUM][2] = {
   /* 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 void get_base_count_mag(int *mag, int *count,
                                       const tran_low_t *tcoeffs, int bwl,
                                       int height, int row, int col) {
@@ -333,9 +330,10 @@ static const int sig_ref_offset_horiz[SIG_REF_OFFSET_NUM][2] = {
 };
 
 #if USE_CAUSAL_BASE_CTX
-static INLINE int get_nz_count_mag(const void *tcoeffs, int bwl, int height,
-                                   int row, int col, TX_CLASS tx_class,
-                                   const int coeff_is_byte_flag, int *mag) {
+static INLINE int get_nz_count_mag(const uint8_t *const levels, const int bwl,
+                                   const int height, const int row,
+                                   const int col, const TX_CLASS tx_class,
+                                   int *const mag) {
   int count = 0;
   *mag = 0;
   for (int idx = 0; idx < SIG_REF_OFFSET_NUM; ++idx) {
@@ -355,8 +353,7 @@ static INLINE int get_nz_count_mag(const void *tcoeffs, int bwl, int height,
         ref_col >= (1 << bwl))
       continue;
     const int nb_pos = (ref_row << bwl) + ref_col;
-    int level = abs(coeff_is_byte_flag ? ((const uint8_t *)tcoeffs)[nb_pos]
-                                       : ((const tran_low_t *)tcoeffs)[nb_pos]);
+    const int level = levels[nb_pos];
     count += (level != 0);
 #if 1
     if (idx < 5) {
@@ -367,9 +364,9 @@ static INLINE int get_nz_count_mag(const void *tcoeffs, int bwl, int height,
   return count;
 }
 #endif
-static INLINE int get_nz_count(const void *tcoeffs, int bwl, int height,
-                               int row, int col, TX_CLASS tx_class,
-                               const int coeff_is_byte_flag) {
+static INLINE int get_nz_count(const uint8_t *const levels, const int bwl,
+                               const int height, const int row, const int col,
+                               const TX_CLASS tx_class) {
   int count = 0;
   for (int idx = 0; idx < SIG_REF_OFFSET_NUM; ++idx) {
     const int ref_row = row + ((tx_class == TX_CLASS_2D)
@@ -386,9 +383,7 @@ static INLINE int get_nz_count(const void *tcoeffs, int bwl, int height,
         ref_col >= (1 << bwl))
       continue;
     const int nb_pos = (ref_row << bwl) + ref_col;
-    count +=
-        ((coeff_is_byte_flag ? ((const uint8_t *)tcoeffs)[nb_pos]
-                             : ((const tran_low_t *)tcoeffs)[nb_pos]) != 0);
+    count += (levels[nb_pos] != 0);
   }
   return count;
 }
@@ -468,10 +463,10 @@ static INLINE int get_nz_map_ctx_from_count(int count,
   }
 }
 
-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) {
+static INLINE int get_nz_map_ctx(const uint8_t *const levels,
+                                 const int scan_idx, const int16_t *const scan,
+                                 const int bwl, const int height,
+                                 const TX_TYPE tx_type) {
   const int coeff_idx = scan[scan_idx];
   const int row = coeff_idx >> bwl;
   const int col = coeff_idx - (row << bwl);
@@ -479,12 +474,10 @@ static INLINE int get_nz_map_ctx(const void *const tcoeffs, const int scan_idx,
   int tx_class = get_tx_class(tx_type);
 #if USE_CAUSAL_BASE_CTX
   int mag = 0;
-  int count = get_nz_count_mag(tcoeffs, bwl, height, row, col, tx_class,
-                               coeff_is_byte_flag, &mag);
+  int count = get_nz_count_mag(levels, bwl, height, row, col, tx_class, &mag);
   return get_nz_map_ctx_from_count(count, coeff_idx, bwl, height, mag, tx_type);
 #else
-  int count = get_nz_count(tcoeffs, bwl, height, row, col, tx_class,
-                           coeff_is_byte_flag);
+  int count = get_nz_count(levels, bwl, height, row, col, tx_class);
   return get_nz_map_ctx_from_count(count, coeff_idx, bwl, height, tx_type);
 #endif
 }

av1/decoder/decodetxb.c

@@ -159,7 +159,7 @@ uint8_t av1_read_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
   for (int i = 0; i < *eob; ++i) {
     c = *eob - 1 - i;
     int is_nz;
-    int coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type, 1);
+    int coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type);
 
     if (c < *eob - 1) {
       is_nz = av1_read_record_bin(

av1/encoder/encodetxb.c

@@ -233,7 +233,7 @@ void get_dist_cost_stats(LevelDownStats *stats, int scan_idx,
   stats->rd_low = RDCOST(txb_info->rdmult, stats->rate_low, stats->dist_low);
 
   int coeff_ctx = get_nz_map_ctx(levels, scan_idx, scan, txb_info->bwl,
-                                 txb_info->height, txb_info->tx_type, 1);
+                                 txb_info->height, txb_info->tx_type);
 
   if ((stats->rd_low < stats->rd) && (stats->low_qc == 0)) {
     stats->nz_rate = txb_costs->nz_map_cost[coeff_ctx][0];
@@ -331,7 +331,7 @@ void av1_write_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
   int coeff_ctx = 0;
   for (int i = 0; i < eob; ++i) {
     c = eob - 1 - i;
-    coeff_ctx = get_nz_map_ctx(tcoeff, c, scan, bwl, height, tx_type, 0);
+    coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type);
 
     tran_low_t v = tcoeff[scan[c]];
     int is_nz = (v != 0);
@@ -356,7 +356,7 @@ void av1_write_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCKD *xd,
 #else
   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, 1);
+    int coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type);
 
     tran_low_t v = tcoeff[scan[c]];
     int is_nz = (v != 0);
@@ -603,12 +603,12 @@ int av1_cost_coeffs_txb(const AV1_COMMON *const cm, MACROBLOCK *x, int plane,
     int is_nz = (v != 0);
     int level = abs(v);
 #if USE_CAUSAL_BASE_CTX
-    coeff_ctx = get_nz_map_ctx(qcoeff, c, scan, bwl, height, tx_type, 0);
+    coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type);
 #endif
 
     if (c < eob - 1) {
 #if !USE_CAUSAL_BASE_CTX
-      int coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type, 1);
+      int coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type);
 #endif
       cost += coeff_costs->nz_map_cost[coeff_ctx][is_nz];
     }
@@ -726,6 +726,7 @@ void gen_txb_cache(TxbCache *txb_cache, TxbInfo *txb_info) {
   const int bwl = txb_info->bwl;
   const int height = txb_info->height;
   const tran_low_t *const qcoeff = txb_info->qcoeff;
+  const uint8_t *const levels = txb_info->levels;
   const BASE_CTX_TABLE *base_ctx_table =
       txb_info->coeff_ctx_table->base_ctx_table;
   for (int c = 0; c < txb_info->eob; ++c) {
@@ -734,7 +735,7 @@ void gen_txb_cache(TxbCache *txb_cache, TxbInfo *txb_info) {
     const int col = coeff_idx - (row << bwl);
 
     txb_cache->nz_count_arr[coeff_idx] = get_nz_count(
-        qcoeff, bwl, height, row, col, get_tx_class(txb_info->tx_type), 0);
+        levels, bwl, height, row, col, get_tx_class(txb_info->tx_type));
 
     const int nz_count = txb_cache->nz_count_arr[coeff_idx];
     txb_cache->nz_ctx_arr[coeff_idx] =
@@ -1314,7 +1315,7 @@ void update_level_down(const int coeff_idx, TxbCache *const txb_cache,
 #endif
             txb_info->tx_type);
         // int ref_ctx = get_nz_map_ctx(txb_info->levels, nb_coeff_idx,
-        // txb_info->bwl, tx_type, 1);
+        // txb_info->bwl, tx_type);
         // 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);
@@ -1406,15 +1407,15 @@ static int get_coeff_cost(const tran_low_t qc, const int scan_idx,
 
 #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);
 #endif
 
   if (scan_idx < txb_info->eob - 1) {
 #if !USE_CAUSAL_BASE_CTX
     int coeff_ctx =
         get_nz_map_ctx(txb_info->levels, scan_idx, scan, txb_info->bwl,
-                       txb_info->height, txb_info->tx_type, 1);
+                       txb_info->height, txb_info->tx_type);
 #endif
     cost += txb_costs->nz_map_cost[coeff_ctx][is_nz];
   }
@@ -2032,7 +2033,7 @@ void av1_update_and_record_txb_context(int plane, int block, int blk_row,
     int is_nz = (v != 0);
 
 #if USE_CAUSAL_BASE_CTX
-    coeff_ctx = get_nz_map_ctx(tcoeff, c, scan, bwl, height, tx_type, 0);
+    coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type);
     if (c < eob - 1) {
       ++(*nz_map_count)[coeff_ctx][is_nz];
       if (allow_update_cdf)
@@ -2054,7 +2055,7 @@ void av1_update_and_record_txb_context(int plane, int block, int blk_row,
       }
     }
 #else
-    int coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type, 1);
+    int coeff_ctx = get_nz_map_ctx(levels, c, scan, bwl, height, tx_type);
 
     if (c == eob - 1) continue;