Use block_idx rather than block_raster_idx

1) block_raster_idx is actually raster order only when tx_size
is TX_4x8.
It's very specific, so we should put it near to the place it's
actually used.

2) Sync the meaning of block_idx on encoder/decoder sides

Change-Id: I7d37a992cb773503e29f9c0d9d2586e580aa6173
diff --git a/av1/common/blockd.h b/av1/common/blockd.h
index 3f40724..8243f2a 100644
--- a/av1/common/blockd.h
+++ b/av1/common/blockd.h
@@ -864,6 +864,27 @@
 #define FIXED_TX_TYPE 0
 #endif
 
+// Converts block_index for given transform size to index of the block in raster
+// order.
+static INLINE int av1_block_index_to_raster_order(TX_SIZE tx_size,
+                                                  int block_idx) {
+  // For transform size 4x8, the possible block_idx values are 0 & 2, because
+  // block_idx values are incremented in steps of size 'tx_width_unit x
+  // tx_height_unit'. But, for this transform size, block_idx = 2 corresponds to
+  // block number 1 in raster order, inside an 8x8 MI block.
+  // For any other transform size, the two indices are equivalent.
+  return (tx_size == TX_4X8 && block_idx == 2) ? 1 : block_idx;
+}
+
+// Inverse of above function.
+// Note: only implemented for transform sizes 4x4, 4x8 and 8x4 right now.
+static INLINE int av1_raster_order_to_block_index(TX_SIZE tx_size,
+                                                  int raster_order) {
+  assert(tx_size == TX_4X4 || tx_size == TX_4X8 || tx_size == TX_8X4);
+  // We ensure that block indices are 0 & 2 if tx size is 4x8 or 8x4.
+  return (tx_size == TX_4X4) ? raster_order : (raster_order > 0) ? 2 : 0;
+}
+
 static INLINE TX_TYPE get_default_tx_type(PLANE_TYPE plane_type,
                                           const MACROBLOCKD *xd, int block_idx,
                                           TX_SIZE tx_size) {
@@ -879,12 +900,14 @@
 }
 
 static INLINE TX_TYPE get_tx_type(PLANE_TYPE plane_type, const MACROBLOCKD *xd,
-                                  int block_idx, TX_SIZE tx_size) {
+                                  int block, TX_SIZE tx_size) {
   const MODE_INFO *const mi = xd->mi[0];
   const MB_MODE_INFO *const mbmi = &mi->mbmi;
 
+  const int block_raster_idx = av1_block_index_to_raster_order(tx_size, block);
+
   if (FIXED_TX_TYPE)
-    return get_default_tx_type(plane_type, xd, block_idx, tx_size);
+    return get_default_tx_type(plane_type, xd, block_raster_idx, tx_size);
 
 #if CONFIG_EXT_TX
   if (xd->lossless[mbmi->segment_id] || txsize_sqr_map[tx_size] > TX_32X32 ||
@@ -921,10 +944,10 @@
     return DCT_DCT;
   else  // Sub8x8 Intra OR UV-Intra
     return intra_mode_to_tx_type_context[plane_type == PLANE_TYPE_Y
-                                             ? get_y_mode(mi, block_idx)
+                                             ? get_y_mode(mi, block_raster_idx)
                                              : mbmi->uv_mode];
 #else   // CONFIG_EXT_TX
-  (void)block_idx;
+  (void)block;
   if (plane_type != PLANE_TYPE_Y || xd->lossless[mbmi->segment_id] ||
       txsize_sqr_map[tx_size] >= TX_32X32)
     return DCT_DCT;
diff --git a/av1/common/reconintra.c b/av1/common/reconintra.c
index 45f3a20..1fee33d 100644
--- a/av1/common/reconintra.c
+++ b/av1/common/reconintra.c
@@ -2256,14 +2256,14 @@
                          plane);
 }
 
-void av1_predict_intra_block_facade(MACROBLOCKD *xd, int plane,
-                                    int block_raster_idx, int blk_col,
-                                    int blk_row, TX_SIZE tx_size) {
+void av1_predict_intra_block_facade(MACROBLOCKD *xd, int plane, int block_idx,
+                                    int blk_col, int blk_row, TX_SIZE tx_size) {
   struct macroblockd_plane *const pd = &xd->plane[plane];
   const int dst_stride = pd->dst.stride;
   uint8_t *dst =
       &pd->dst.buf[(blk_row * dst_stride + blk_col) << tx_size_wide_log2[0]];
   MB_MODE_INFO *mbmi = &xd->mi[0]->mbmi;
+  int block_raster_idx = av1_block_index_to_raster_order(tx_size, block_idx);
   PREDICTION_MODE mode =
       (plane == 0) ? get_y_mode(xd->mi[0], block_raster_idx) : mbmi->uv_mode;
   av1_predict_intra_block(xd, pd->width, pd->height, txsize_to_bsize[tx_size],
diff --git a/av1/common/reconintra.h b/av1/common/reconintra.h
index 4d3ab6b..7ee0c49 100644
--- a/av1/common/reconintra.h
+++ b/av1/common/reconintra.h
@@ -20,9 +20,8 @@
 #endif
 
 void av1_init_intra_predictors(void);
-void av1_predict_intra_block_facade(MACROBLOCKD *xd, int plane,
-                                    int block_raster_idx, int blk_col,
-                                    int blk_row, TX_SIZE tx_size);
+void av1_predict_intra_block_facade(MACROBLOCKD *xd, int plane, int block_idx,
+                                    int blk_col, int blk_row, TX_SIZE tx_size);
 void av1_predict_intra_block(const MACROBLOCKD *xd, int bw, int bh,
                              BLOCK_SIZE bsize, PREDICTION_MODE mode,
                              const uint8_t *ref, int ref_stride, uint8_t *dst,
diff --git a/av1/decoder/decodeframe.c b/av1/decoder/decodeframe.c
index 5fad99f..e3f45cf 100644
--- a/av1/decoder/decodeframe.c
+++ b/av1/decoder/decodeframe.c
@@ -504,11 +504,31 @@
 }
 #endif
 
+static int get_block_idx(const MACROBLOCKD *xd, int plane, int row, int col) {
+  const int bsize = xd->mi[0]->mbmi.sb_type;
+  const struct macroblockd_plane *pd = &xd->plane[plane];
+#if CONFIG_CB4X4
+#if CONFIG_CHROMA_2X2
+  const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, pd);
+#else
+  const BLOCK_SIZE plane_bsize =
+      AOMMAX(BLOCK_4X4, get_plane_block_size(bsize, pd));
+#endif  // CONFIG_CHROMA_2X2
+#else
+  const BLOCK_SIZE plane_bsize =
+      get_plane_block_size(AOMMAX(BLOCK_8X8, bsize), pd);
+#endif
+  const int max_blocks_wide = max_block_wide(xd, plane_bsize, plane);
+  const TX_SIZE tx_size = get_tx_size(plane, xd);
+  const uint8_t txh_unit = tx_size_high_unit[tx_size];
+  return row * max_blocks_wide + col * txh_unit;
+}
+
 static void predict_and_reconstruct_intra_block(
     AV1_COMMON *cm, MACROBLOCKD *const xd, aom_reader *const r,
     MB_MODE_INFO *const mbmi, int plane, int row, int col, TX_SIZE tx_size) {
   PLANE_TYPE plane_type = get_plane_type(plane);
-  const int block_idx = (row << 1) + col;
+  const int block_idx = get_block_idx(xd, plane, row, col);
 #if CONFIG_PVQ
   (void)r;
 #endif
@@ -563,7 +583,7 @@
 
   if (tx_size == plane_tx_size) {
     PLANE_TYPE plane_type = get_plane_type(plane);
-    int block_idx = (blk_row << 1) + blk_col;
+    int block_idx = get_block_idx(xd, plane, blk_row, blk_col);
     TX_TYPE tx_type = get_tx_type(plane_type, xd, block_idx, plane_tx_size);
 #if CONFIG_LV_MAP
     (void)segment_id;
@@ -610,7 +630,7 @@
                                    int plane, int row, int col,
                                    TX_SIZE tx_size) {
   PLANE_TYPE plane_type = get_plane_type(plane);
-  int block_idx = (row << 1) + col;
+  int block_idx = get_block_idx(xd, plane, row, col);
   TX_TYPE tx_type = get_tx_type(plane_type, xd, block_idx, tx_size);
 #if CONFIG_PVQ
   int eob;
diff --git a/av1/encoder/block.h b/av1/encoder/block.h
index 28b7e10..a4ecdc2 100644
--- a/av1/encoder/block.h
+++ b/av1/encoder/block.h
@@ -226,27 +226,6 @@
 #endif
 };
 
-// Converts block_index for given transform size to index of the block in raster
-// order.
-static INLINE int av1_block_index_to_raster_order(TX_SIZE tx_size,
-                                                  int block_idx) {
-  // For transform size 4x8, the possible block_idx values are 0 & 2, because
-  // block_idx values are incremented in steps of size 'tx_width_unit x
-  // tx_height_unit'. But, for this transform size, block_idx = 2 corresponds to
-  // block number 1 in raster order, inside an 8x8 MI block.
-  // For any other transform size, the two indices are equivalent.
-  return (tx_size == TX_4X8 && block_idx == 2) ? 1 : block_idx;
-}
-
-// Inverse of above function.
-// Note: only implemented for transform sizes 4x4, 4x8 and 8x4 right now.
-static INLINE int av1_raster_order_to_block_index(TX_SIZE tx_size,
-                                                  int raster_order) {
-  assert(tx_size == TX_4X4 || tx_size == TX_4X8 || tx_size == TX_8X4);
-  // We ensure that block indices are 0 & 2 if tx size is 4x8 or 8x4.
-  return (tx_size == TX_4X4) ? raster_order : (raster_order > 0) ? 2 : 0;
-}
-
 #ifdef __cplusplus
 }  // extern "C"
 #endif
diff --git a/av1/encoder/encodemb.c b/av1/encoder/encodemb.c
index 02b7ab0..1df0dc3 100644
--- a/av1/encoder/encodemb.c
+++ b/av1/encoder/encodemb.c
@@ -159,8 +159,7 @@
   const int default_eob = tx_size_2d[tx_size];
   const int16_t *const dequant_ptr = pd->dequant;
   const uint8_t *const band_translate = get_band_translate(tx_size);
-  const int block_raster_idx = av1_block_index_to_raster_order(tx_size, block);
-  TX_TYPE tx_type = get_tx_type(plane_type, xd, block_raster_idx, tx_size);
+  TX_TYPE tx_type = get_tx_type(plane_type, xd, block, tx_size);
   const SCAN_ORDER *const scan_order =
       get_scan(cm, tx_size, tx_type, is_inter_block(&xd->mi[0]->mbmi));
   const int16_t *const scan = scan_order->scan;
@@ -554,8 +553,7 @@
   struct macroblockd_plane *const pd = &xd->plane[plane];
 #endif
   PLANE_TYPE plane_type = get_plane_type(plane);
-  const int block_raster_idx = av1_block_index_to_raster_order(tx_size, block);
-  TX_TYPE tx_type = get_tx_type(plane_type, xd, block_raster_idx, tx_size);
+  TX_TYPE tx_type = get_tx_type(plane_type, xd, block, tx_size);
   const int is_inter = is_inter_block(mbmi);
   const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, is_inter);
   tran_low_t *const coeff = BLOCK_OFFSET(p->coeff, block);
@@ -749,7 +747,6 @@
 #if !CONFIG_PVQ
   ENTROPY_CONTEXT *a, *l;
 #endif
-  const int block_raster_idx = av1_block_index_to_raster_order(tx_size, block);
 #if CONFIG_VAR_TX
   int bw = block_size_wide[plane_bsize] >> tx_size_wide_log2[0];
 #endif
@@ -800,7 +797,7 @@
 
   if (x->pvq_skip[plane]) return;
 #endif
-  TX_TYPE tx_type = get_tx_type(pd->plane_type, xd, block_raster_idx, tx_size);
+  TX_TYPE tx_type = get_tx_type(pd->plane_type, xd, block, tx_size);
   av1_inverse_transform_block(xd, dqcoeff, tx_type, tx_size, dst,
                               pd->dst.stride, p->eobs[block]);
 }
@@ -1069,16 +1066,12 @@
   struct macroblockd_plane *const pd = &xd->plane[plane];
   tran_low_t *dqcoeff = BLOCK_OFFSET(pd->dqcoeff, block);
   PLANE_TYPE plane_type = get_plane_type(plane);
-  const int block_raster_idx = av1_block_index_to_raster_order(tx_size, block);
-  const TX_TYPE tx_type =
-      get_tx_type(plane_type, xd, block_raster_idx, tx_size);
+  const TX_TYPE tx_type = get_tx_type(plane_type, xd, block, tx_size);
   uint16_t *eob = &p->eobs[block];
   const int dst_stride = pd->dst.stride;
   uint8_t *dst =
       &pd->dst.buf[(blk_row * dst_stride + blk_col) << tx_size_wide_log2[0]];
-
-  av1_predict_intra_block_facade(xd, plane, block_raster_idx, blk_col, blk_row,
-                                 tx_size);
+  av1_predict_intra_block_facade(xd, plane, block, blk_col, blk_row, tx_size);
   av1_subtract_txb(x, plane, plane_bsize, blk_col, blk_row, tx_size);
 
   const ENTROPY_CONTEXT *a = &args->ta[blk_col];
diff --git a/av1/encoder/rdopt.c b/av1/encoder/rdopt.c
index a76eb09..3fa9302 100644
--- a/av1/encoder/rdopt.c
+++ b/av1/encoder/rdopt.c
@@ -1524,11 +1524,8 @@
         (void)dst;
 #endif  // !CONFIG_PVQ
 
-        const int block_raster_idx =
-            av1_block_index_to_raster_order(tx_size, block);
         const PLANE_TYPE plane_type = get_plane_type(plane);
-        TX_TYPE tx_type =
-            get_tx_type(plane_type, xd, block_raster_idx, tx_size);
+        TX_TYPE tx_type = get_tx_type(plane_type, xd, block, tx_size);
 
         av1_inverse_transform_block(xd, dqcoeff, tx_type, tx_size, recon,
                                     MAX_TX_SIZE, eob);
@@ -1574,15 +1571,12 @@
                                            *(args->t_left + blk_row));
   RD_STATS this_rd_stats;
 
-  const int block_raster_idx = av1_block_index_to_raster_order(tx_size, block);
-
   av1_init_rd_stats(&this_rd_stats);
 
   if (args->exit_early) return;
 
   if (!is_inter_block(mbmi)) {
-    av1_predict_intra_block_facade(xd, plane, block_raster_idx, blk_col,
-                                   blk_row, tx_size);
+    av1_predict_intra_block_facade(xd, plane, block, blk_col, blk_row, tx_size);
     av1_subtract_txb(x, plane, plane_bsize, blk_col, blk_row, tx_size);
   }
 
@@ -1597,8 +1591,7 @@
     struct macroblockd_plane *const pd = &xd->plane[plane];
     tran_low_t *dqcoeff = BLOCK_OFFSET(pd->dqcoeff, block);
     PLANE_TYPE plane_type = get_plane_type(plane);
-    const TX_TYPE tx_type =
-        get_tx_type(plane_type, xd, block_raster_idx, tx_size);
+    const TX_TYPE tx_type = get_tx_type(plane_type, xd, block, tx_size);
     const int dst_stride = pd->dst.stride;
     uint8_t *dst =
         &pd->dst.buf[(blk_row * dst_stride + blk_col) << tx_size_wide_log2[0]];
@@ -2359,10 +2352,7 @@
   int block = 0;
   for (row = 0; row < max_blocks_high; row += stepr) {
     for (col = 0; col < max_blocks_wide; col += stepc) {
-      const int block_raster_idx =
-          av1_block_index_to_raster_order(tx_size, block);
-      av1_predict_intra_block_facade(xd, 0, block_raster_idx, col, row,
-                                     tx_size);
+      av1_predict_intra_block_facade(xd, 0, block, col, row, tx_size);
       block += step;
     }
   }
@@ -2701,8 +2691,7 @@
                                     src_stride, dst, dst_stride, xd->bd);
 #endif
           if (is_lossless) {
-            TX_TYPE tx_type =
-                get_tx_type(PLANE_TYPE_Y, xd, block_raster_idx, tx_size);
+            TX_TYPE tx_type = get_tx_type(PLANE_TYPE_Y, xd, block, tx_size);
             const SCAN_ORDER *scan_order = get_scan(cm, tx_size, tx_type, 0);
             const int coeff_ctx =
                 combine_entropy_contexts(tempa[idx], templ[idy]);
@@ -2746,8 +2735,7 @@
           } else {
             int64_t dist;
             unsigned int tmp;
-            TX_TYPE tx_type =
-                get_tx_type(PLANE_TYPE_Y, xd, block_raster_idx, tx_size);
+            TX_TYPE tx_type = get_tx_type(PLANE_TYPE_Y, xd, block, tx_size);
             const SCAN_ORDER *scan_order = get_scan(cm, tx_size, tx_type, 0);
             const int coeff_ctx =
                 combine_entropy_contexts(tempa[idx], templ[idy]);
@@ -2897,8 +2885,7 @@
 #endif  // !CONFIG_PVQ
 
         if (is_lossless) {
-          TX_TYPE tx_type =
-              get_tx_type(PLANE_TYPE_Y, xd, block_raster_idx, tx_size);
+          TX_TYPE tx_type = get_tx_type(PLANE_TYPE_Y, xd, block, tx_size);
           const SCAN_ORDER *scan_order = get_scan(cm, tx_size, tx_type, 0);
           const int coeff_ctx =
               combine_entropy_contexts(tempa[idx], templ[idy]);
@@ -2955,8 +2942,7 @@
         } else {
           int64_t dist;
           unsigned int tmp;
-          TX_TYPE tx_type =
-              get_tx_type(PLANE_TYPE_Y, xd, block_raster_idx, tx_size);
+          TX_TYPE tx_type = get_tx_type(PLANE_TYPE_Y, xd, block, tx_size);
           const SCAN_ORDER *scan_order = get_scan(cm, tx_size, tx_type, 0);
           const int coeff_ctx =
               combine_entropy_contexts(tempa[idx], templ[idy]);
diff --git a/av1/encoder/tokenize.c b/av1/encoder/tokenize.c
index 2a354ab..bbc5a97 100644
--- a/av1/encoder/tokenize.c
+++ b/av1/encoder/tokenize.c
@@ -288,8 +288,7 @@
   struct macroblockd_plane *pd = &xd->plane[plane];
   const PLANE_TYPE type = pd->plane_type;
   const int ref = is_inter_block(mbmi);
-  const int block_raster_idx = av1_block_index_to_raster_order(tx_size, block);
-  const TX_TYPE tx_type = get_tx_type(type, xd, block_raster_idx, tx_size);
+  const TX_TYPE tx_type = get_tx_type(type, xd, block, tx_size);
   const SCAN_ORDER *const scan_order = get_scan(cm, tx_size, tx_type, ref);
   const int rate = av1_cost_coeffs(cm, x, plane, block, tx_size, scan_order,
                                    pd->above_context + blk_col,
@@ -465,8 +464,7 @@
   const int segment_id = mbmi->segment_id;
 #endif  // CONFIG_SUEPRTX
   const int16_t *scan, *nb;
-  const int block_raster_idx = av1_block_index_to_raster_order(tx_size, block);
-  const TX_TYPE tx_type = get_tx_type(type, xd, block_raster_idx, tx_size);
+  const TX_TYPE tx_type = get_tx_type(type, xd, block, tx_size);
   const SCAN_ORDER *const scan_order =
       get_scan(cm, tx_size, tx_type, is_inter_block(mbmi));
   const int ref = is_inter_block(mbmi);