Refactor av1_set_contexts() to support 2x2 transform block unit

Replace hard coded numbers with table fetch. Support scalability in
transform block operating unit.

Change-Id: I4a5e67576e72de4936ac52d1401412ca6217ed6d

av1/common/blockd.c

@@ -164,17 +164,19 @@ void av1_foreach_transformed_block(const MACROBLOCKD *const xd,
 
 #if !CONFIG_PVQ
 void av1_set_contexts(const MACROBLOCKD *xd, struct macroblockd_plane *pd,
-                      TX_SIZE tx_size, int has_eob, int aoff, int loff) {
+                      int plane, TX_SIZE tx_size, int has_eob, int aoff,
+                      int loff) {
   ENTROPY_CONTEXT *const a = pd->above_context + aoff;
   ENTROPY_CONTEXT *const l = pd->left_context + loff;
   const int txs_wide = tx_size_wide_unit[tx_size];
   const int txs_high = tx_size_high_unit[tx_size];
+  const BLOCK_SIZE bsize = AOMMAX(xd->mi[0]->mbmi.sb_type, BLOCK_8X8);
+  const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, pd);
 
   // above
   if (has_eob && xd->mb_to_right_edge < 0) {
     int i;
-    const int blocks_wide =
-        pd->n4_w + (xd->mb_to_right_edge >> (5 + pd->subsampling_x));
+    const int blocks_wide = max_block_wide(xd, plane_bsize, plane);
     int above_contexts = txs_wide;
     if (above_contexts + aoff > blocks_wide)
       above_contexts = blocks_wide - aoff;
@@ -188,8 +190,7 @@ void av1_set_contexts(const MACROBLOCKD *xd, struct macroblockd_plane *pd,
   // left
   if (has_eob && xd->mb_to_bottom_edge < 0) {
     int i;
-    const int blocks_high =
-        pd->n4_h + (xd->mb_to_bottom_edge >> (5 + pd->subsampling_y));
+    const int blocks_high = max_block_high(xd, plane_bsize, plane);
     int left_contexts = txs_high;
     if (left_contexts + loff > blocks_high) left_contexts = blocks_high - loff;
 

av1/common/blockd.h

@@ -879,7 +879,8 @@ void av1_foreach_transformed_block_interleave(
 #endif
 
 void av1_set_contexts(const MACROBLOCKD *xd, struct macroblockd_plane *pd,
-                      TX_SIZE tx_size, int has_eob, int aoff, int loff);
+                      int plane, TX_SIZE tx_size, int has_eob, int aoff,
+                      int loff);
 
 #if CONFIG_EXT_INTER
 static INLINE int is_interintra_allowed_bsize(const BLOCK_SIZE bsize) {

av1/decoder/detokenize.c

@@ -365,7 +365,7 @@ int av1_decode_block_tokens(MACROBLOCKD *const xd, int plane,
 #endif  // CONFIG_NEW_QUANT
                    ctx, sc->scan, sc->neighbors, max_scan_line, r);
 #endif  // CONFIG_AOM_QM
-  av1_set_contexts(xd, pd, tx_size, eob > 0, x, y);
+  av1_set_contexts(xd, pd, plane, tx_size, eob > 0, x, y);
   return eob;
 }
 #endif

av1/encoder/tokenize.c

@@ -341,7 +341,8 @@ static void cost_coeffs_b(int plane, int block, int blk_row, int blk_col,
                              scan_order->neighbors, 0);
   args->this_rate += rate;
   (void)plane_bsize;
-  av1_set_contexts(xd, pd, tx_size, p->eobs[block] > 0, blk_col, blk_row);
+  av1_set_contexts(xd, pd, plane, tx_size, p->eobs[block] > 0, blk_col,
+                   blk_row);
 }
 
 static void set_entropy_context_b(int plane, int block, int blk_row,
@@ -354,7 +355,8 @@ static void set_entropy_context_b(int plane, int block, int blk_row,
   struct macroblock_plane *p = &x->plane[plane];
   struct macroblockd_plane *pd = &xd->plane[plane];
   (void)plane_bsize;
-  av1_set_contexts(xd, pd, tx_size, p->eobs[block] > 0, blk_col, blk_row);
+  av1_set_contexts(xd, pd, plane, tx_size, p->eobs[block] > 0, blk_col,
+                   blk_row);
 }
 
 static INLINE void add_token(TOKENEXTRA **t, const aom_prob *context_tree,
@@ -518,7 +520,7 @@ static void tokenize_b(int plane, int block, int blk_row, int blk_col,
   av1_update_scan_count_facade((AV1_COMMON *)cm, tx_size, tx_type, qcoeff, c);
 #endif
 
-  av1_set_contexts(xd, pd, tx_size, c > 0, blk_col, blk_row);
+  av1_set_contexts(xd, pd, plane, tx_size, c > 0, blk_col, blk_row);
 }
 #endif