Replace stride by bwl in get_nz_count

Change-Id: Iaa464fdd36c9a97070517757645da68ef009b82c

av1/common/txb_common.h

diff --git a/av1/common/txb_common.h b/av1/common/txb_common.h
index 72de63e..8cec13d 100644
--- a/av1/common/txb_common.h
+++ b/av1/common/txb_common.h
@@ -195,17 +195,17 @@
   { -1, 1 },  { 0, -2 }, { 0, -1 }, { 1, -2 },  { 1, -1 },
 };
 
-static INLINE int get_nz_count(const tran_low_t *tcoeffs, int stride,
-                               int height, int row, int col,
-                               const int16_t *iscan) {
+static INLINE int get_nz_count(const tran_low_t *tcoeffs, int bwl, int height,
+                               int row, int col, const int16_t *iscan) {
   int count = 0;
-  const int pos = row * stride + col;
+  const int pos = (row << bwl) + col;
   for (int idx = 0; idx < SIG_REF_OFFSET_NUM; ++idx) {
     const int ref_row = row + sig_ref_offset[idx][0];
     const int ref_col = col + sig_ref_offset[idx][1];
-    if (ref_row < 0 || ref_col < 0 || ref_row >= height || ref_col >= stride)
+    if (ref_row < 0 || ref_col < 0 || ref_row >= height ||
+        ref_col >= (1 << bwl))
       continue;
-    const int nb_pos = ref_row * stride + ref_col;
+    const int nb_pos = (ref_row << bwl) + ref_col;
     if (iscan[nb_pos] < iscan[pos]) count += (tcoeffs[nb_pos] != 0);
   }
   return count;
@@ -267,10 +267,9 @@
                                  const int coeff_idx,  // raster order
                                  const int bwl, const int height,
                                  const int16_t *iscan) {
-  int stride = 1 << bwl;
   const int row = coeff_idx >> bwl;
   const int col = coeff_idx - (row << bwl);
-  int count = get_nz_count(tcoeffs, stride, height, row, col, iscan);
+  int count = get_nz_count(tcoeffs, bwl, height, row, col, iscan);
   return get_nz_map_ctx_from_count(count, tcoeffs, coeff_idx, bwl, iscan);
 }
 

av1/encoder/encodetxb.c

diff --git a/av1/encoder/encodetxb.c b/av1/encoder/encodetxb.c
index 6a28a6a..c09d1f7 100644
--- a/av1/encoder/encodetxb.c
+++ b/av1/encoder/encodetxb.c
@@ -431,16 +431,16 @@
 }
 
 static void gen_nz_count_arr(int(*nz_count_arr), const tran_low_t *qcoeff,
-                             int stride, int height, int eob,
+                             int bwl, int height, int eob,
                              const SCAN_ORDER *scan_order) {
   const int16_t *scan = scan_order->scan;
   const int16_t *iscan = scan_order->iscan;
   for (int c = 0; c < eob; ++c) {
     const int coeff_idx = scan[c];  // raster order
-    const int row = coeff_idx / stride;
-    const int col = coeff_idx % stride;
+    const int row = coeff_idx >> bwl;
+    const int col = coeff_idx - (row << bwl);
     nz_count_arr[coeff_idx] =
-        get_nz_count(qcoeff, stride, height, row, col, iscan);
+        get_nz_count(qcoeff, bwl, height, row, col, iscan);
   }
 }
 
@@ -547,7 +547,7 @@
 // TODO(angiebird): add static once this function is called
 void gen_txb_cache(TxbCache *txb_cache, TxbInfo *txb_info) {
   const int16_t *scan = txb_info->scan_order->scan;
-  gen_nz_count_arr(txb_cache->nz_count_arr, txb_info->qcoeff, txb_info->stride,
+  gen_nz_count_arr(txb_cache->nz_count_arr, txb_info->qcoeff, txb_info->bwl,
                    txb_info->height, txb_info->eob, txb_info->scan_order);
   gen_nz_ctx_arr(txb_cache->nz_ctx_arr, txb_cache->nz_count_arr,
                  txb_info->qcoeff, txb_info->bwl, txb_info->eob,